wehub-resource-sync
231 lines
7.0 KiB
Python
231 lines
7.0 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import json
|
|
import sys
|
|
from types import SimpleNamespace
|
|
from typing import Any
|
|
|
|
import pytest
|
|
|
|
from opensquilla.cli.gateway_client import (
|
|
GatewayClient,
|
|
_normalize_session_error_payload,
|
|
_task_terminal_as_session_event,
|
|
)
|
|
|
|
_STOP = object()
|
|
|
|
|
|
class _FakeWebSocket:
|
|
def __init__(self, recv_frames: list[dict[str, Any]] | None = None) -> None:
|
|
self._recv_frames = list(recv_frames or [])
|
|
self.sent: list[str] = []
|
|
self.iter_queue: asyncio.Queue[str | object] = asyncio.Queue()
|
|
self.closed = False
|
|
|
|
async def recv(self) -> str:
|
|
if not self._recv_frames:
|
|
raise AssertionError("unexpected recv() call")
|
|
return json.dumps(self._recv_frames.pop(0))
|
|
|
|
async def send(self, payload: str) -> None:
|
|
self.sent.append(payload)
|
|
|
|
async def close(self) -> None:
|
|
self.closed = True
|
|
|
|
def __aiter__(self) -> _FakeWebSocket:
|
|
return self
|
|
|
|
async def __anext__(self) -> str:
|
|
item = await self.iter_queue.get()
|
|
if item is _STOP:
|
|
raise StopAsyncIteration
|
|
assert isinstance(item, str)
|
|
return item
|
|
|
|
|
|
class _BrokenSendWebSocket:
|
|
async def send(self, payload: str) -> None:
|
|
raise RuntimeError("socket already closed")
|
|
|
|
|
|
async def _wait_for(predicate, *, timeout: float = 1.0) -> None:
|
|
deadline = asyncio.get_running_loop().time() + timeout
|
|
while not predicate():
|
|
if asyncio.get_running_loop().time() >= deadline:
|
|
raise AssertionError("condition was not met before timeout")
|
|
await asyncio.sleep(0.005)
|
|
|
|
|
|
def _install_fake_websockets(monkeypatch: pytest.MonkeyPatch, ws: _FakeWebSocket) -> None:
|
|
async def _connect(url: str) -> _FakeWebSocket:
|
|
return ws
|
|
|
|
monkeypatch.setitem(sys.modules, "websockets", SimpleNamespace(connect=_connect))
|
|
|
|
|
|
def test_task_terminal_event_uses_terminal_message() -> None:
|
|
event = _task_terminal_as_session_event(
|
|
"task.timeout",
|
|
{
|
|
"task_id": "task-1",
|
|
"terminal_reason": "timeout",
|
|
"terminal_message": "The task timed out before it could finish.",
|
|
"error_message": "Gateway task timeout: Stream idle for more than 60s",
|
|
},
|
|
)
|
|
|
|
assert event is not None
|
|
assert event["event"] == "session.event.error"
|
|
assert event["message"] == "The task timed out before it could finish."
|
|
assert "Gateway task" not in event["message"]
|
|
|
|
|
|
def test_session_error_payload_is_normalized_for_gateway_client() -> None:
|
|
payload = _normalize_session_error_payload(
|
|
{
|
|
"message": "Iteration 1 exceeded iteration_timeout",
|
|
"code": "iteration_timeout",
|
|
}
|
|
)
|
|
|
|
assert payload["message"] == "The task timed out before it could finish."
|
|
assert payload["terminal_message"] == "The task timed out before it could finish."
|
|
assert payload["error_message"] == "The task timed out before it could finish."
|
|
|
|
|
|
def test_gateway_client_error_normalization_preserves_generic_error_detail() -> None:
|
|
payload = _normalize_session_error_payload(
|
|
{
|
|
"message": "Tool failed with exit code 2",
|
|
"code": "agent_error",
|
|
}
|
|
)
|
|
|
|
assert payload["message"] == "The task failed before it could finish."
|
|
assert payload["error_message"] == "Tool failed with exit code 2"
|
|
|
|
|
|
def _handshake_frames(*, keepalive_ms: int = 60_000) -> list[dict[str, Any]]:
|
|
return [
|
|
{"type": "event", "event": "connect.challenge", "payload": {"nonce": "n"}},
|
|
{
|
|
"type": "hello-ok",
|
|
"policy": {"client_ws_keepalive_timeout_ms": keepalive_ms},
|
|
},
|
|
]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_heartbeat_interval_derived_from_hello_policy(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
ws = _FakeWebSocket(_handshake_frames(keepalive_ms=60_000))
|
|
_install_fake_websockets(monkeypatch, ws)
|
|
client = GatewayClient()
|
|
|
|
await client.connect()
|
|
try:
|
|
assert client._heartbeat_interval == 24.0 # noqa: SLF001
|
|
assert client._heartbeat_task is not None # noqa: SLF001
|
|
finally:
|
|
await client.close()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_heartbeat_interval_stays_below_short_hello_policy(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
ws = _FakeWebSocket(_handshake_frames(keepalive_ms=2_000))
|
|
_install_fake_websockets(monkeypatch, ws)
|
|
client = GatewayClient()
|
|
|
|
await client.connect()
|
|
try:
|
|
assert client._heartbeat_interval == 0.8 # noqa: SLF001
|
|
finally:
|
|
await client.close()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_heartbeat_loop_sends_text_ping_not_rpc() -> None:
|
|
ws = _FakeWebSocket()
|
|
client = GatewayClient()
|
|
client._ws = ws # noqa: SLF001
|
|
client._heartbeat_interval = 0.01 # noqa: SLF001
|
|
|
|
task = asyncio.create_task(client._heartbeat_loop()) # noqa: SLF001
|
|
try:
|
|
await _wait_for(lambda: bool(ws.sent))
|
|
finally:
|
|
task.cancel()
|
|
with pytest.raises(asyncio.CancelledError):
|
|
await task
|
|
|
|
frame = json.loads(ws.sent[0])
|
|
assert frame == {"type": "ping"}
|
|
assert "id" not in frame
|
|
assert "method" not in frame
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_listener_ignores_pong_frames() -> None:
|
|
ws = _FakeWebSocket()
|
|
client = GatewayClient()
|
|
client._ws = ws # noqa: SLF001
|
|
|
|
task = asyncio.create_task(client._listen()) # noqa: SLF001
|
|
try:
|
|
await ws.iter_queue.put(json.dumps({"type": "pong"}))
|
|
await asyncio.sleep(0.02)
|
|
assert client._recv_queue.empty() # noqa: SLF001
|
|
finally:
|
|
task.cancel()
|
|
with pytest.raises(asyncio.CancelledError):
|
|
await task
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_close_cancels_heartbeat(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
ws = _FakeWebSocket(_handshake_frames(keepalive_ms=60_000))
|
|
_install_fake_websockets(monkeypatch, ws)
|
|
client = GatewayClient()
|
|
|
|
await client.connect()
|
|
heartbeat_task = client._heartbeat_task # noqa: SLF001
|
|
assert heartbeat_task is not None
|
|
|
|
await client.close()
|
|
|
|
assert heartbeat_task.done()
|
|
assert ws.closed is True
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_listener_close_stops_heartbeat_task() -> None:
|
|
ws = _FakeWebSocket()
|
|
client = GatewayClient()
|
|
client._ws = ws # noqa: SLF001
|
|
client._heartbeat_interval = 60.0 # noqa: SLF001
|
|
client._heartbeat_task = asyncio.create_task(client._heartbeat_loop(ws)) # noqa: SLF001
|
|
|
|
listener_task = asyncio.create_task(client._listen()) # noqa: SLF001
|
|
await ws.iter_queue.put(_STOP)
|
|
await _wait_for(lambda: client._connection_error is not None) # noqa: SLF001
|
|
|
|
assert client._heartbeat_task.done() # noqa: SLF001
|
|
await listener_task
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_call_after_send_failure_raises_clear_connection_error() -> None:
|
|
client = GatewayClient()
|
|
client._ws = _BrokenSendWebSocket() # noqa: SLF001
|
|
|
|
with pytest.raises(ConnectionError, match="Gateway connection lost"):
|
|
await client._call("sessions.messages.subscribe", {"key": "agent:main:x"}) # noqa: SLF001
|
|
|
|
assert client._pending == {} # noqa: SLF001
|