wehub-resource-sync
259 lines
8.3 KiB
Python
259 lines
8.3 KiB
Python
"""Tests for the in-memory host connection registry."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from dataclasses import dataclass, field
|
|
|
|
import pytest
|
|
|
|
from omnigent.host.frames import HostHelloFrame
|
|
from omnigent.server.host_registry import HostRegistry, RunnerExitReports
|
|
|
|
|
|
@dataclass
|
|
class FakeWebSocket:
|
|
"""Minimal WebSocket fake for registry tests.
|
|
|
|
Records sent text frames so tests can assert on outbound
|
|
messages without a real network.
|
|
|
|
:param sent: List of text frames sent via ``send_text``.
|
|
"""
|
|
|
|
sent: list[str] = field(default_factory=list)
|
|
|
|
async def send_text(self, data: str) -> None:
|
|
"""Record a sent frame.
|
|
|
|
:param data: The text frame content.
|
|
"""
|
|
self.sent.append(data)
|
|
|
|
async def receive_text(self) -> str:
|
|
"""Block forever — tests don't read from the fake.
|
|
|
|
:returns: Never returns in practice.
|
|
"""
|
|
await asyncio.sleep(3600)
|
|
return "" # pragma: no cover
|
|
|
|
|
|
def _make_hello(name: str = "test-host") -> HostHelloFrame:
|
|
"""Build a minimal HostHelloFrame for tests.
|
|
|
|
:param name: Human-readable host name.
|
|
:returns: A :class:`HostHelloFrame` with default values.
|
|
"""
|
|
return HostHelloFrame(
|
|
version="0.1.0",
|
|
frame_protocol_version=1,
|
|
name=name,
|
|
)
|
|
|
|
|
|
def test_register_and_get() -> None:
|
|
"""
|
|
Verify that a registered host can be retrieved by ID.
|
|
|
|
If get() returns None after register(), the registry's
|
|
internal dict is not being populated.
|
|
"""
|
|
registry = HostRegistry()
|
|
ws = FakeWebSocket()
|
|
conn = registry.register("host_aaa", ws, _make_hello(), owner="alice")
|
|
|
|
fetched = registry.get("host_aaa")
|
|
assert fetched is conn
|
|
assert fetched is not None
|
|
assert fetched.host_id == "host_aaa"
|
|
assert fetched.owner == "alice"
|
|
assert fetched.hello.name == "test-host"
|
|
|
|
|
|
def test_deregister() -> None:
|
|
"""
|
|
Verify that deregister removes the host from the registry.
|
|
|
|
If get() still returns the connection after deregister(), the
|
|
pop() call is missing or targeting the wrong key.
|
|
"""
|
|
registry = HostRegistry()
|
|
registry.register("host_bbb", FakeWebSocket(), _make_hello(), owner="bob")
|
|
registry.deregister("host_bbb")
|
|
|
|
assert registry.get("host_bbb") is None
|
|
|
|
|
|
def test_deregister_noop_for_unknown() -> None:
|
|
"""
|
|
Verify that deregister is a no-op for an unknown host_id.
|
|
|
|
The disconnect callback may fire for hosts that failed
|
|
registration; it must not raise.
|
|
"""
|
|
registry = HostRegistry()
|
|
registry.deregister("host_nonexistent")
|
|
|
|
|
|
def test_online_host_ids() -> None:
|
|
"""
|
|
Verify that online_host_ids returns all registered hosts
|
|
and updates after deregister.
|
|
|
|
If a deregistered host still appears, the dict pop is broken.
|
|
If a registered host is missing, the dict insert is broken.
|
|
"""
|
|
registry = HostRegistry()
|
|
registry.register("host_c1", FakeWebSocket(), _make_hello(), owner="carol")
|
|
registry.register("host_c2", FakeWebSocket(), _make_hello(), owner="carol")
|
|
|
|
ids = registry.online_host_ids()
|
|
assert set(ids) == {"host_c1", "host_c2"}
|
|
|
|
registry.deregister("host_c1")
|
|
ids = registry.online_host_ids()
|
|
assert ids == ["host_c2"]
|
|
|
|
|
|
def test_register_replaces_stale_connection() -> None:
|
|
"""
|
|
Verify that registering the same host_id replaces the old
|
|
connection (newest wins) and poisons the old outbound queue.
|
|
|
|
If the old connection isn't replaced, a reconnecting host would
|
|
have two live entries, causing frame routing confusion.
|
|
"""
|
|
registry = HostRegistry()
|
|
old_ws = FakeWebSocket()
|
|
old_conn = registry.register("host_ddd", old_ws, _make_hello(), owner="dave")
|
|
|
|
new_ws = FakeWebSocket()
|
|
new_conn = registry.register("host_ddd", new_ws, _make_hello(), owner="dave")
|
|
|
|
# New connection is the one returned by get().
|
|
assert registry.get("host_ddd") is new_conn
|
|
assert new_conn is not old_conn
|
|
|
|
# Old connection's outbound queue was poisoned with None.
|
|
# The None sentinel tells the sender loop to exit.
|
|
poison = old_conn.outbound_queue.get_nowait()
|
|
assert poison is None
|
|
|
|
|
|
def test_send_text_enqueues_frame() -> None:
|
|
"""
|
|
Verify that send_text puts the frame on the connection's
|
|
outbound queue.
|
|
|
|
If the frame doesn't appear in the queue, the sender loop
|
|
would never transmit it and the host would never receive the
|
|
launch request.
|
|
"""
|
|
registry = HostRegistry()
|
|
ws = FakeWebSocket()
|
|
conn = registry.register("host_eee", ws, _make_hello(), owner="eve")
|
|
|
|
registry.send_text(conn, '{"kind": "host.launch_runner"}')
|
|
|
|
# Frame should be on the outbound queue.
|
|
frame = conn.outbound_queue.get_nowait()
|
|
assert frame == '{"kind": "host.launch_runner"}'
|
|
|
|
|
|
def test_send_text_raises_if_connection_replaced() -> None:
|
|
"""
|
|
Verify that send_text raises ConnectionError if the connection
|
|
was replaced by a newer one.
|
|
|
|
Without this guard, a stale reference could enqueue frames on
|
|
a dead connection's queue, which would never be drained.
|
|
"""
|
|
registry = HostRegistry()
|
|
old_conn = registry.register("host_fff", FakeWebSocket(), _make_hello(), owner="frank")
|
|
registry.register("host_fff", FakeWebSocket(), _make_hello(), owner="frank")
|
|
|
|
with pytest.raises(ConnectionError, match="connection was replaced"):
|
|
registry.send_text(old_conn, '{"kind": "test"}')
|
|
|
|
|
|
def test_get_returns_none_for_unknown() -> None:
|
|
"""
|
|
Verify that get() returns None for a host that was never
|
|
registered.
|
|
"""
|
|
registry = HostRegistry()
|
|
assert registry.get("host_nonexistent") is None
|
|
|
|
|
|
# ── RunnerExitReports ───────────────────────────────────
|
|
|
|
|
|
def test_exit_reports_record_and_get_visible_for_owner() -> None:
|
|
"""
|
|
A recorded exit report is readable by the owner with its exact
|
|
error text.
|
|
|
|
The error message is the entire diagnostic value of the report —
|
|
if it comes back mangled or None, the waiting client falls back
|
|
to the blind 60s timeout this feature removes.
|
|
"""
|
|
reports = RunnerExitReports()
|
|
reports.record("runner_abc", "runner process exited with code 1", "alice")
|
|
|
|
assert reports.get_visible("runner_abc", "alice") == ("runner process exited with code 1")
|
|
|
|
|
|
def test_exit_reports_hidden_from_other_users() -> None:
|
|
"""
|
|
Another user's report reads as None (W6-2 posture).
|
|
|
|
The report's log tail can contain agent output, so a non-owner
|
|
must see nothing — the same "reveal nothing about other users'
|
|
runners" rule the status endpoint applies to ``online``.
|
|
"""
|
|
reports = RunnerExitReports()
|
|
reports.record("runner_abc", "runner process exited with code 1", "alice")
|
|
|
|
assert reports.get_visible("runner_abc", "bob") is None
|
|
|
|
|
|
def test_exit_reports_visible_when_auth_disabled() -> None:
|
|
"""
|
|
With auth disabled on both sides (owner None, user None) the
|
|
report is readable — single-user/local mode must not lose the
|
|
diagnostic.
|
|
"""
|
|
reports = RunnerExitReports()
|
|
reports.record("runner_abc", "runner process exited with code 1", None)
|
|
|
|
assert reports.get_visible("runner_abc", None) == ("runner process exited with code 1")
|
|
|
|
|
|
def test_exit_reports_missing_runner_returns_none() -> None:
|
|
"""
|
|
An unknown runner id reads as None — the status endpoint then
|
|
omits the ``error`` field rather than inventing one.
|
|
"""
|
|
reports = RunnerExitReports()
|
|
assert reports.get_visible("runner_never_reported", "alice") is None
|
|
|
|
|
|
def test_exit_reports_get_is_unscoped() -> None:
|
|
"""
|
|
The unscoped ``get`` returns the error regardless of owner.
|
|
|
|
Used by the session snapshot, which has already authorized access
|
|
by session permission — the report is that session's own runner, so
|
|
no owner re-check is needed (unlike the auth-less status endpoint,
|
|
which must use ``get_visible``). If ``get`` ever started scoping by
|
|
owner, the snapshot would silently drop the error for shared
|
|
sessions viewed by a non-owner.
|
|
"""
|
|
reports = RunnerExitReports()
|
|
reports.record("runner_abc", "runner process exited with code 1", "alice")
|
|
|
|
assert reports.get("runner_abc") == "runner process exited with code 1"
|
|
# Missing runner still reads None (snapshot then leaves last_task_error unset).
|
|
assert reports.get("runner_unknown") is None
|