622 lines
19 KiB
Python
622 lines
19 KiB
Python
"""Unit tests for WebSocket Hub — connection lifecycle, routing, and edge cases.
|
|
|
|
Uses unittest.mock for WebSocket mocking, no live server needed.
|
|
Run: python -m pytest tests/test_hub_unit.py -v
|
|
"""
|
|
|
|
import asyncio
|
|
import json
|
|
import time
|
|
import pytest
|
|
import pytest_asyncio
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
|
|
# Add project root to sys.path
|
|
import sys
|
|
import os
|
|
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
|
|
|
# Configure pytest-asyncio
|
|
pytestmark = pytest.mark.anyio
|
|
|
|
from auth import TokenManager
|
|
from hub import WSHub, WSConnection, MsgType, PER_CONN_RATE_LIMIT
|
|
|
|
|
|
# ─── Fixtures ───
|
|
|
|
def make_token_manager(secret="test-secret-64chars-aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa",
|
|
reg_code="test-reg-code"):
|
|
return TokenManager(secret=secret, registration_code=reg_code)
|
|
|
|
|
|
def make_mock_ws(closed=False):
|
|
"""Create a mock WebSocket with async send_json."""
|
|
ws = AsyncMock()
|
|
ws.closed = closed
|
|
ws.send_json = AsyncMock()
|
|
ws.close = AsyncMock()
|
|
return ws
|
|
|
|
|
|
def make_connection(conn_id, project="test_project", pc_name="test_pc",
|
|
instance_number=1, authenticated=True, ws=None):
|
|
"""Create a WSConnection with a mock WebSocket."""
|
|
if ws is None:
|
|
ws = make_mock_ws()
|
|
conn = WSConnection(
|
|
conn_id=conn_id,
|
|
ws=ws,
|
|
project=project,
|
|
pc_name=pc_name,
|
|
instance_number=instance_number,
|
|
authenticated=authenticated,
|
|
)
|
|
return conn
|
|
|
|
|
|
def make_hub():
|
|
"""Create a Hub with a TokenManager (no auth enforcement)."""
|
|
tm = make_token_manager()
|
|
return WSHub(tm)
|
|
|
|
|
|
def register_conn(hub, conn):
|
|
"""Manually register a connection in the hub (skip auth flow)."""
|
|
hub.connections[conn.conn_id] = conn
|
|
if conn.project not in hub.project_connections:
|
|
hub.project_connections[conn.project] = set()
|
|
hub.project_connections[conn.project].add(conn.conn_id)
|
|
|
|
|
|
# ─── Connection Tracking ───
|
|
|
|
class TestConnectionTracking:
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_register_connection_adds_to_tracking(self):
|
|
hub = make_hub()
|
|
conn = make_connection("c1")
|
|
hub._register_connection(conn)
|
|
|
|
assert "c1" in hub.connections
|
|
assert "c1" in hub.project_connections["test_project"]
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_register_multiple_connections_same_project(self):
|
|
hub = make_hub()
|
|
c1 = make_connection("c1")
|
|
c2 = make_connection("c2")
|
|
hub._register_connection(c1)
|
|
hub._register_connection(c2)
|
|
|
|
assert hub.get_active_count("test_project") == 2
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_register_connections_different_projects(self):
|
|
hub = make_hub()
|
|
c1 = make_connection("c1", project="proj_a")
|
|
c2 = make_connection("c2", project="proj_b")
|
|
hub._register_connection(c1)
|
|
hub._register_connection(c2)
|
|
|
|
assert hub.get_active_count("proj_a") == 1
|
|
assert hub.get_active_count("proj_b") == 1
|
|
assert len(hub.connections) == 2
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_disconnect_removes_from_tracking(self):
|
|
hub = make_hub()
|
|
conn = make_connection("c1")
|
|
register_conn(hub, conn)
|
|
|
|
await hub._disconnect(conn)
|
|
|
|
assert "c1" not in hub.connections
|
|
assert hub.get_active_count("test_project") == 0
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_disconnect_removes_empty_project(self):
|
|
hub = make_hub()
|
|
conn = make_connection("c1")
|
|
register_conn(hub, conn)
|
|
|
|
await hub._disconnect(conn)
|
|
|
|
assert "test_project" not in hub.project_connections
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_disconnect_keeps_other_connections(self):
|
|
hub = make_hub()
|
|
c1 = make_connection("c1")
|
|
c2 = make_connection("c2")
|
|
register_conn(hub, c1)
|
|
register_conn(hub, c2)
|
|
|
|
await hub._disconnect(c1)
|
|
|
|
assert "c1" not in hub.connections
|
|
assert "c2" in hub.connections
|
|
assert hub.get_active_count("test_project") == 1
|
|
|
|
|
|
# ─── Instance Number Assignment ───
|
|
|
|
class TestInstanceNumbers:
|
|
|
|
def test_first_instance_gets_1(self):
|
|
hub = make_hub()
|
|
num = hub._assign_instance_number("proj", "c1")
|
|
assert num == 1
|
|
|
|
def test_second_instance_gets_2(self):
|
|
hub = make_hub()
|
|
c1 = make_connection("c1", instance_number=1)
|
|
register_conn(hub, c1)
|
|
|
|
num = hub._assign_instance_number("proj", "c2")
|
|
# project "proj" has no connections (c1 is in "test_project"), so it gets 1
|
|
assert num == 1
|
|
|
|
def test_gap_filling(self):
|
|
"""If instance #1 disconnects, next connection gets #1 again."""
|
|
hub = make_hub()
|
|
c2 = make_connection("c2", instance_number=2)
|
|
register_conn(hub, c2)
|
|
|
|
num = hub._assign_instance_number("test_project", "c3")
|
|
assert num == 1 # #1 is available
|
|
|
|
def test_sequential_assignment(self):
|
|
hub = make_hub()
|
|
c1 = make_connection("c1", instance_number=1)
|
|
c2 = make_connection("c2", instance_number=2)
|
|
register_conn(hub, c1)
|
|
register_conn(hub, c2)
|
|
|
|
num = hub._assign_instance_number("test_project", "c3")
|
|
assert num == 3
|
|
|
|
|
|
# ─── Pending Owners Lifecycle ───
|
|
|
|
class TestPendingOwners:
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_pending_tracks_owner(self):
|
|
"""Processing a pending message tracks the owner connection."""
|
|
hub = make_hub()
|
|
conn = make_connection("c1")
|
|
register_conn(hub, conn)
|
|
|
|
await hub._handle_message(conn, MsgType.PENDING, {
|
|
"data": {"request_id": "req-001", "command": "ls"}
|
|
})
|
|
|
|
assert hub.pending_owners["req-001"] == "c1"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_auto_resolve_clears_owner(self):
|
|
"""Auto-resolve removes the pending owner tracking."""
|
|
hub = make_hub()
|
|
conn = make_connection("c1")
|
|
register_conn(hub, conn)
|
|
hub.pending_owners["req-001"] = "c1"
|
|
|
|
await hub._handle_message(conn, MsgType.AUTO_RESOLVE, {
|
|
"data": {"request_id": "req-001"}
|
|
})
|
|
|
|
assert "req-001" not in hub.pending_owners
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_disconnect_reassigns_pending_to_same_project(self):
|
|
"""When owner disconnects, pending is reassigned to another connection
|
|
in the same project (not deleted)."""
|
|
hub = make_hub()
|
|
c1 = make_connection("c1")
|
|
c2 = make_connection("c2")
|
|
register_conn(hub, c1)
|
|
register_conn(hub, c2)
|
|
hub.pending_owners["req-001"] = "c1"
|
|
|
|
await hub._disconnect(c1)
|
|
|
|
# req-001 should now be owned by c2
|
|
assert hub.pending_owners.get("req-001") == "c2"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_disconnect_deletes_pending_when_no_remaining(self):
|
|
"""When owner disconnects and no other connections exist, pending is deleted."""
|
|
hub = make_hub()
|
|
conn = make_connection("c1")
|
|
register_conn(hub, conn)
|
|
hub.pending_owners["req-001"] = "c1"
|
|
|
|
await hub._disconnect(conn)
|
|
|
|
assert "req-001" not in hub.pending_owners
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_register_reassigns_orphaned_pending(self):
|
|
"""When a new connection registers, orphaned pending_owners
|
|
(pointing to dead conn_ids) are reassigned to it."""
|
|
hub = make_hub()
|
|
# Simulate orphaned pending (old conn_id "dead" not in connections)
|
|
hub.pending_owners["req-001"] = "dead-conn"
|
|
hub.pending_owners["req-002"] = "dead-conn"
|
|
|
|
new_conn = make_connection("c-new")
|
|
hub._register_connection(new_conn)
|
|
|
|
assert hub.pending_owners["req-001"] == "c-new"
|
|
assert hub.pending_owners["req-002"] == "c-new"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_disconnect_reassigns_only_to_authenticated(self):
|
|
"""Pending should only be reassigned to authenticated connections."""
|
|
hub = make_hub()
|
|
c1 = make_connection("c1")
|
|
c2 = make_connection("c2", authenticated=False) # Not authenticated
|
|
register_conn(hub, c1)
|
|
register_conn(hub, c2)
|
|
hub.pending_owners["req-001"] = "c1"
|
|
|
|
await hub._disconnect(c1)
|
|
|
|
# c2 is not authenticated, so pending should be deleted
|
|
assert "req-001" not in hub.pending_owners
|
|
|
|
|
|
# ─── Response Routing ───
|
|
|
|
class TestResponseRouting:
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_send_response_to_owner(self):
|
|
hub = make_hub()
|
|
conn = make_connection("c1")
|
|
register_conn(hub, conn)
|
|
hub.pending_owners["req-001"] = "c1"
|
|
|
|
result = await hub.send_response_to_pending_owner(
|
|
"req-001", {"type": "response", "approved": True}
|
|
)
|
|
|
|
assert result is True
|
|
# Owner should be cleaned up after successful delivery
|
|
assert "req-001" not in hub.pending_owners
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_send_response_owner_dead_fallback(self):
|
|
"""When original owner is dead, response routes to another connection
|
|
in the same project."""
|
|
hub = make_hub()
|
|
dead_conn = make_connection("c-dead", ws=make_mock_ws(closed=True))
|
|
alive_conn = make_connection("c-alive")
|
|
register_conn(hub, dead_conn)
|
|
register_conn(hub, alive_conn)
|
|
hub.pending_owners["req-001"] = "c-dead"
|
|
|
|
result = await hub.send_response_to_pending_owner(
|
|
"req-001", {"type": "response", "approved": True}
|
|
)
|
|
|
|
assert result is True # Rerouted successfully
|
|
assert "req-001" not in hub.pending_owners
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_send_response_no_owner(self):
|
|
"""When no owner exists for a request_id, return False."""
|
|
hub = make_hub()
|
|
|
|
result = await hub.send_response_to_pending_owner(
|
|
"req-unknown", {"type": "response"}
|
|
)
|
|
|
|
assert result is False
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_send_response_all_dead(self):
|
|
"""When owner and all project connections are dead, return False."""
|
|
hub = make_hub()
|
|
dead = make_connection("c-dead", ws=make_mock_ws(closed=True))
|
|
register_conn(hub, dead)
|
|
hub.pending_owners["req-001"] = "c-dead"
|
|
|
|
result = await hub.send_response_to_pending_owner(
|
|
"req-001", {"type": "response"}
|
|
)
|
|
|
|
assert result is False
|
|
assert "req-001" not in hub.pending_owners # Cleaned up
|
|
|
|
|
|
# ─── Message Routing ───
|
|
|
|
class TestMessageRouting:
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_broadcast_to_project(self):
|
|
hub = make_hub()
|
|
c1 = make_connection("c1")
|
|
c2 = make_connection("c2")
|
|
c3 = make_connection("c3", project="other_project")
|
|
register_conn(hub, c1)
|
|
register_conn(hub, c2)
|
|
register_conn(hub, c3)
|
|
|
|
msg = {"type": "test", "data": "hello"}
|
|
await hub.broadcast_to_project("test_project", msg)
|
|
|
|
# c1 and c2 should have received the message (via queue)
|
|
assert c1.send_queue.qsize() == 1
|
|
assert c2.send_queue.qsize() == 1
|
|
assert c3.send_queue.qsize() == 0 # Different project
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_send_to_instance(self):
|
|
hub = make_hub()
|
|
c1 = make_connection("c1", instance_number=1)
|
|
c2 = make_connection("c2", instance_number=2)
|
|
register_conn(hub, c1)
|
|
register_conn(hub, c2)
|
|
|
|
result = await hub.send_to_instance("test_project", 2, {"type": "cmd"})
|
|
|
|
assert result is True
|
|
assert c1.send_queue.qsize() == 0
|
|
assert c2.send_queue.qsize() == 1
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_send_to_instance_not_found(self):
|
|
hub = make_hub()
|
|
c1 = make_connection("c1", instance_number=1)
|
|
register_conn(hub, c1)
|
|
|
|
result = await hub.send_to_instance("test_project", 99, {"type": "cmd"})
|
|
|
|
assert result is False
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_broadcast_skips_unauthenticated(self):
|
|
hub = make_hub()
|
|
c1 = make_connection("c1", authenticated=True)
|
|
c2 = make_connection("c2", authenticated=False)
|
|
register_conn(hub, c1)
|
|
register_conn(hub, c2)
|
|
|
|
await hub.broadcast_to_project("test_project", {"type": "msg"})
|
|
|
|
assert c1.send_queue.qsize() == 1
|
|
assert c2.send_queue.qsize() == 0
|
|
|
|
|
|
# ─── Message Handling (Dispatch) ───
|
|
|
|
class TestMessageHandling:
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_pending_sets_metadata(self):
|
|
"""Pending messages should have source metadata injected."""
|
|
hub = make_hub()
|
|
received = []
|
|
hub.set_bot_handlers(on_pending=AsyncMock(side_effect=lambda p, d: received.append(d)))
|
|
|
|
conn = make_connection("c1", pc_name="my_pc", instance_number=3)
|
|
register_conn(hub, conn)
|
|
|
|
await hub._handle_message(conn, MsgType.PENDING, {
|
|
"data": {"request_id": "req-x", "command": "npm install"}
|
|
})
|
|
|
|
assert len(received) == 1
|
|
payload = received[0]
|
|
assert payload["_conn_id"] == "c1"
|
|
assert payload["_instance_number"] == 3
|
|
assert payload["_pc_name"] == "my_pc"
|
|
assert payload["project_name"] == "test_project"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_chat_sets_project_name(self):
|
|
hub = make_hub()
|
|
received = []
|
|
hub.set_bot_handlers(on_chat=AsyncMock(side_effect=lambda p, d: received.append(d)))
|
|
|
|
conn = make_connection("c1", project="my_proj")
|
|
register_conn(hub, conn)
|
|
|
|
await hub._handle_message(conn, MsgType.CHAT, {
|
|
"data": {"content": "hello"}
|
|
})
|
|
|
|
assert received[0]["project_name"] == "my_proj"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_no_handler_doesnt_crash(self):
|
|
"""If no bot handlers are set, messages should be silently dropped."""
|
|
hub = make_hub()
|
|
conn = make_connection("c1")
|
|
register_conn(hub, conn)
|
|
|
|
# Should not raise
|
|
await hub._handle_message(conn, MsgType.PENDING, {"data": {"request_id": "r1"}})
|
|
await hub._handle_message(conn, MsgType.CHAT, {"data": {"content": "hi"}})
|
|
await hub._handle_message(conn, MsgType.REGISTER, {"data": {}})
|
|
await hub._handle_message(conn, MsgType.AUTO_RESOLVE, {"data": {}})
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_heartbeat_updates_timestamp(self):
|
|
hub = make_hub()
|
|
conn = make_connection("c1")
|
|
register_conn(hub, conn)
|
|
old_ts = conn.last_heartbeat
|
|
|
|
# Small delay to ensure timestamp difference
|
|
await asyncio.sleep(0.01)
|
|
await hub._handle_message(conn, MsgType.HEARTBEAT, {})
|
|
|
|
assert conn.last_heartbeat > old_ts
|
|
|
|
|
|
# ─── Rate Limiting ───
|
|
|
|
class TestRateLimiting:
|
|
|
|
def test_under_limit_allows(self):
|
|
hub = make_hub()
|
|
conn = make_connection("c1")
|
|
|
|
for _ in range(PER_CONN_RATE_LIMIT - 1):
|
|
assert hub._check_rate_limit(conn) is True
|
|
|
|
def test_at_limit_blocks(self):
|
|
hub = make_hub()
|
|
conn = make_connection("c1")
|
|
|
|
# Fill up to limit
|
|
for _ in range(PER_CONN_RATE_LIMIT):
|
|
hub._check_rate_limit(conn)
|
|
|
|
# Next should be blocked
|
|
assert hub._check_rate_limit(conn) is False
|
|
|
|
def test_old_timestamps_expire(self):
|
|
hub = make_hub()
|
|
conn = make_connection("c1")
|
|
|
|
# Add old timestamps (beyond RATE_WINDOW)
|
|
old_time = time.time() - 20.0
|
|
conn._msg_timestamps = [old_time] * PER_CONN_RATE_LIMIT
|
|
|
|
# Should be allowed (old timestamps are pruned)
|
|
assert hub._check_rate_limit(conn) is True
|
|
|
|
|
|
# ─── Deduplication ───
|
|
|
|
class TestDeduplication:
|
|
|
|
def test_first_message_not_duplicate(self):
|
|
hub = make_hub()
|
|
assert hub._is_duplicate("msg-001") is False
|
|
|
|
def test_same_id_is_duplicate(self):
|
|
hub = make_hub()
|
|
hub._is_duplicate("msg-001")
|
|
assert hub._is_duplicate("msg-001") is True
|
|
|
|
def test_different_id_not_duplicate(self):
|
|
hub = make_hub()
|
|
hub._is_duplicate("msg-001")
|
|
assert hub._is_duplicate("msg-002") is False
|
|
|
|
def test_cleanup_on_overflow(self):
|
|
hub = make_hub()
|
|
# Add many entries to trigger cleanup
|
|
for i in range(10001):
|
|
hub._recent_msg_ids[f"msg-{i}"] = time.time()
|
|
|
|
# Should trigger cleanup and still work
|
|
hub._is_duplicate("new-msg")
|
|
assert len(hub._recent_msg_ids) <= 10002 # +1 for new-msg
|
|
|
|
|
|
# ─── Queue Backpressure ───
|
|
|
|
class TestQueueBackpressure:
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_queue_send_normal(self):
|
|
hub = make_hub()
|
|
conn = make_connection("c1")
|
|
|
|
await hub._queue_send(conn, {"type": "test"})
|
|
|
|
assert conn.send_queue.qsize() == 1
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_queue_full_drops_oldest(self):
|
|
hub = make_hub()
|
|
conn = make_connection("c1")
|
|
|
|
# Fill the queue
|
|
for i in range(100):
|
|
await hub._queue_send(conn, {"type": "test", "i": i})
|
|
|
|
# Queue should be full (100 items)
|
|
assert conn.send_queue.qsize() == 100
|
|
|
|
# Adding one more should drop the oldest
|
|
await hub._queue_send(conn, {"type": "test", "i": "new"})
|
|
|
|
assert conn.send_queue.qsize() == 100 # Still 100
|
|
|
|
|
|
# ─── Diagnostics ───
|
|
|
|
class TestDiagnostics:
|
|
|
|
def test_get_status_empty(self):
|
|
hub = make_hub()
|
|
status = hub.get_status()
|
|
|
|
assert status["total_connections"] == 0
|
|
assert status["projects"] == {}
|
|
assert status["pending_owners"] == 0
|
|
|
|
def test_get_status_with_connections(self):
|
|
hub = make_hub()
|
|
c1 = make_connection("c1", project="proj_a", instance_number=1)
|
|
c2 = make_connection("c2", project="proj_a", instance_number=2)
|
|
register_conn(hub, c1)
|
|
register_conn(hub, c2)
|
|
hub.pending_owners["req-001"] = "c1"
|
|
|
|
status = hub.get_status()
|
|
|
|
assert status["total_connections"] == 2
|
|
assert status["projects"]["proj_a"]["count"] == 2
|
|
assert status["pending_owners"] == 1
|
|
assert len(status["projects"]["proj_a"]["instances"]) == 2
|
|
|
|
|
|
# ─── Auth (TokenManager) ───
|
|
|
|
class TestTokenManager:
|
|
|
|
def test_create_and_verify_token(self):
|
|
tm = make_token_manager()
|
|
token = tm.create_token("my_project", "my_pc")
|
|
payload = tm.verify_token(token)
|
|
|
|
assert payload is not None
|
|
assert payload["project"] == "my_project"
|
|
assert payload["pc"] == "my_pc"
|
|
|
|
def test_expired_token_rejected(self):
|
|
tm = make_token_manager()
|
|
token = tm.create_token("proj", "pc", ttl=-1) # Already expired
|
|
assert tm.verify_token(token) is None
|
|
|
|
def test_tampered_token_rejected(self):
|
|
tm = make_token_manager()
|
|
token = tm.create_token("proj", "pc")
|
|
tampered = token[:-1] + ("A" if token[-1] != "A" else "B")
|
|
assert tm.verify_token(tampered) is None
|
|
|
|
def test_wrong_secret_rejected(self):
|
|
tm1 = make_token_manager(secret="secret-1" + "a" * 56)
|
|
tm2 = make_token_manager(secret="secret-2" + "b" * 56)
|
|
token = tm1.create_token("proj", "pc")
|
|
assert tm2.verify_token(token) is None
|
|
|
|
def test_registration_code_valid(self):
|
|
tm = make_token_manager(reg_code="my-code")
|
|
assert tm.validate_registration_code("my-code") is True
|
|
assert tm.validate_registration_code("wrong") is False
|
|
|
|
def test_registration_code_empty_allows_all(self):
|
|
tm = make_token_manager(reg_code="")
|
|
assert tm.validate_registration_code("anything") is True
|