summaryrefslogtreecommitdiff
path: root/test_balancing.py
blob: 0f7933decd6cbf5715b03a496b202de9dc9c8d97 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
#!/usr/bin/env python3
"""Real-world load balancing test for vpnem API."""
import urllib.request
import json
import sys

BASE = "https://vpn.em-sysadmin.xyz"

def req(path, method="GET", data=None, headers=None):
    """Make HTTP request."""
    url = BASE + path
    h = headers or {}
    h["Content-Type"] = "application/json"
    body = json.dumps(data).encode() if data else None
    rq = urllib.request.Request(url, data=body, headers=h, method=method)
    try:
        resp = urllib.request.urlopen(rq, timeout=10)
        return json.loads(resp.read().decode())
    except urllib.error.HTTPError as e:
        print(f"  ERROR {e.code}: {e.read().decode()[:200]}")
        return None
    except Exception as e:
        print(f"  ERROR: {e}")
        return None

def get_recommend(client_ip):
    """Get recommendation for a client IP."""
    return req("/api/v1/recommend", headers={"X-Forwarded-For": client_ip})

def connect(client_ip, server_ip, node_id):
    """Report connection."""
    return req("/api/v1/connect", method="POST", data={
        "server_ip": server_ip, "node_id": node_id
    }, headers={"X-Forwarded-For": client_ip})

def disconnect(client_ip):
    """Report disconnection."""
    return req("/api/v1/disconnect", method="POST", data={
        "server_ip": ""
    }, headers={"X-Forwarded-For": client_ip})

def main():
    print("=" * 70)
    print("VPNEM Load Balancing — Real Server Test")
    print("=" * 70)
    print()

    # Step 1: Test recommendations for multiple "studios"
    print("[1] Testing recommendations for 5 different studios...")
    print()
    studios = [
        ("195.10.20.1", "Barnaul Studio 1"),
        ("195.10.20.2", "Barnaul Studio 2"),
        ("91.50.60.1",  "Moscow Studio"),
        ("46.30.20.1",  "Novosibirsk Studio"),
        ("178.120.1.1", "Test Studio"),
    ]

    recommendations = {}
    for ip, name in studios:
        rec = get_recommend(ip)
        if rec:
            server = rec.get("recommended_server_ip", "none")
            reason = rec.get("reason", "unknown")
            load = rec.get("load_info", "no info")
            recommendations[ip] = server
            print(f"  {name:25s} → {server:15s} ({reason})")
            print(f"                            Load: {load}")
        else:
            print(f"  {name:25s} → ERROR")

    print()

    # Step 2: Check distribution — do different studios get different servers?
    print("[2] Distribution check...")
    servers_used = list(recommendations.values())
    unique = set(servers_used)
    print(f"  Studios: {len(studios)}, Recommended IPs: {len(unique)}")
    print(f"  Servers used: {', '.join(unique)}")
    if len(unique) > 1:
        print("  ✅ GOOD — different studios get different servers")
    else:
        print("  ⚠️  All studios got the same server (possible if load is equal)")

    print()

    # Step 3: Simulate connections
    print("[3] Simulating connections...")
    nodes = {
        "5.180.97.181": "nl-multi-181",
        "5.180.97.197": "nl-multi-197",
        "5.180.97.198": "nl-multi-198",
        "5.180.97.199": "nl-multi-199",
    }
    for ip, name in studios:
        srv = recommendations.get(ip)
        nid = nodes.get(srv, "nl-multi-181")
        resp = connect(ip, srv, nid)
        if resp:
            rec_srv = resp.get("recommended_server_ip", "none")
            print(f"  {name:25s} connected → {srv:15s}  (next rec: {rec_srv})")
        else:
            print(f"  {name:25s} connect FAILED")

    print()

    # Step 4: Check recommendations AFTER connections (should shift)
    print("[4] Recommendations AFTER connections (should re-balance)...")
    new_ip = "100.100.100.1"
    rec = get_recommend(new_ip)
    if rec:
        srv = rec.get("recommended_server_ip", "none")
        reason = rec.get("reason", "unknown")
        load = rec.get("load_info", "no info")
        print(f"  New Studio ({new_ip}) → {srv:15s} ({reason})")
        print(f"                            Load: {load}")
        # The new studio should NOT get the most loaded server
    else:
        print(f"  ERROR getting recommendation")

    print()

    # Step 5: Disconnect one studio, check if recommendation changes
    print("[5] Disconnecting first studio and checking rebalance...")
    disc_ip = studios[0][0]
    disconnect(disc_ip)
    print(f"  Disconnected {studios[0][1]} ({disc_ip})")

    rec = get_recommend("200.200.200.1")
    if rec:
        srv = rec.get("recommended_server_ip", "none")
        load = rec.get("load_info", "no info")
        print(f"  New Studio (200.200.200.1) → {srv:15s}")
        print(f"                            Load: {load}")
    else:
        print("  ERROR")

    print()

    # Step 6: Clean up
    print("[6] Cleaning up — disconnecting all test studios...")
    for ip, name in studios:
        disconnect(ip)
    disconnect(new_ip)
    disconnect("200.200.200.1")
    print("  Done")

    print()
    print("=" * 70)
    print("Test complete.")
    print("=" * 70)

if __name__ == "__main__":
    main()