forked from Zakaria/hermes-agent
Zakaria
129 lines
4.3 KiB
Python
129 lines
4.3 KiB
Python
import asyncio
|
|
import threading
|
|
import time
|
|
|
|
from tui_gateway import server
|
|
from tui_gateway import ws as ws_mod
|
|
|
|
|
|
def _run_disconnect(monkeypatch, seed):
|
|
"""Drive handle_ws to its disconnect `finally`, seeding sessions against the
|
|
live WSTransport the moment it exists. Returns nothing; inspect _sessions."""
|
|
# Disable the grace-reap Timer: detached sessions normally schedule a
|
|
# threading.Timer via _schedule_ws_orphan_reap, which would outlive the test
|
|
# and fire _reap during interpreter teardown — touching _sessions/DB and
|
|
# producing spurious post-run errors under the per-file CI runner. Grace=0
|
|
# short-circuits the Timer (see _schedule_ws_orphan_reap) so the test leaves
|
|
# no lingering thread.
|
|
monkeypatch.setattr(server, "_WS_ORPHAN_REAP_GRACE_S", 0)
|
|
|
|
# Mirror the real _finalize_session chokepoint: it is the single place that
|
|
# closes the slash-worker (#38095). Stub it but keep that behavior so the
|
|
# disconnect-reap path still exercises worker teardown.
|
|
def _fake_finalize(s, end_reason="tui_close"):
|
|
w = s.get("slash_worker")
|
|
if w:
|
|
w.close()
|
|
|
|
monkeypatch.setattr(server, "_finalize_session", _fake_finalize)
|
|
|
|
created = []
|
|
real_transport = ws_mod.WSTransport
|
|
monkeypatch.setattr(
|
|
ws_mod, "WSTransport",
|
|
lambda ws, loop, **kw: created.append(real_transport(ws, loop, **kw)) or created[-1],
|
|
)
|
|
|
|
class FakeWS:
|
|
async def accept(self):
|
|
pass
|
|
|
|
async def send_text(self, line):
|
|
pass
|
|
|
|
async def receive_text(self):
|
|
seed(created[0]) # transport now exists; attach it to sessions
|
|
raise ws_mod._WebSocketDisconnect()
|
|
|
|
async def close(self):
|
|
pass
|
|
|
|
asyncio.run(ws_mod.handle_ws(FakeWS()))
|
|
|
|
|
|
def test_ws_disconnect_reaps_flagged_session_and_closes_worker(monkeypatch):
|
|
closed = []
|
|
|
|
class FakeWorker:
|
|
def close(self):
|
|
closed.append(True)
|
|
|
|
server._sessions.clear()
|
|
try:
|
|
_run_disconnect(
|
|
monkeypatch,
|
|
lambda t: server._sessions.update(
|
|
flagged={
|
|
"transport": t,
|
|
"close_on_disconnect": True,
|
|
"slash_worker": FakeWorker(),
|
|
"session_key": "k",
|
|
}
|
|
),
|
|
)
|
|
assert "flagged" not in server._sessions
|
|
assert closed == [True]
|
|
finally:
|
|
server._sessions.clear()
|
|
|
|
|
|
def test_ws_disconnect_preserves_and_repoints_reconnectable_session(monkeypatch):
|
|
server._sessions.clear()
|
|
try:
|
|
_run_disconnect(
|
|
monkeypatch,
|
|
lambda t: server._sessions.update(
|
|
plain={"transport": t, "close_on_disconnect": False, "session_key": "k"}
|
|
),
|
|
)
|
|
assert server._sessions["plain"]["transport"] is server._detached_ws_transport
|
|
finally:
|
|
server._sessions.clear()
|
|
|
|
|
|
def test_ws_write_loop_stall_does_not_latch_transport(monkeypatch):
|
|
"""A write that times out because the event loop is stalled (GIL-heavy
|
|
agent turn) must NOT latch the transport closed — the frame is already
|
|
scheduled and flushes when the loop recovers. Latching here permanently
|
|
silenced live watch windows after one slow write."""
|
|
monkeypatch.setattr(ws_mod, "_WS_WRITE_TIMEOUT_S", 0.05)
|
|
sent = []
|
|
|
|
class FakeWS:
|
|
async def send_text(self, line):
|
|
sent.append(line)
|
|
|
|
loop = asyncio.new_event_loop()
|
|
thread = threading.Thread(target=loop.run_forever, daemon=True)
|
|
thread.start()
|
|
try:
|
|
transport = ws_mod.WSTransport(FakeWS(), loop, peer="stall-test")
|
|
# Stall the loop well past the write timeout, then write from this
|
|
# (non-loop) thread: the wait times out but the send stays in flight.
|
|
loop.call_soon_threadsafe(time.sleep, 0.3)
|
|
assert transport.write({"a": 1}) is True
|
|
assert transport._closed is False
|
|
|
|
# Once the loop breathes again, both the stalled frame and new writes
|
|
# must reach the socket.
|
|
assert transport.write({"b": 2}) is True
|
|
deadline = time.time() + 2
|
|
while len(sent) < 2 and time.time() < deadline:
|
|
time.sleep(0.01)
|
|
assert len(sent) == 2
|
|
assert transport._closed is False
|
|
finally:
|
|
loop.call_soon_threadsafe(loop.stop)
|
|
thread.join(timeout=2)
|
|
loop.close()
|