@samitouri / QOSami_Agents_Zero
Code Commits Issues Pulls Sessions CI/CD
main
main
py 117 lines 3.76 KB
Raw
1 import asyncio
2 import contextlib
3 import socket
4 from typing import Any, AsyncIterator
5
6 import pytest
7
8
9 @contextlib.asynccontextmanager
10 async def _run_asgi_app(app: Any) -> AsyncIterator[str]:
11 import uvicorn
12
13 sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
14 sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
15 sock.bind(("127.0.0.1", 0))
16 sock.listen(128)
17
18 port = sock.getsockname()[1]
19
20 config = uvicorn.Config(
21 app,
22 host="127.0.0.1",
23 port=port,
24 log_level="warning",
25 access_log=False,
26 lifespan="off",
27 )
28 server = uvicorn.Server(config)
29 server.install_signal_handlers = lambda: None # type: ignore[method-assign]
30
31 task = asyncio.create_task(server.serve(sockets=[sock]))
32 try:
33 while not server.started:
34 await asyncio.sleep(0.01)
35 yield f"http://127.0.0.1:{port}"
36 finally:
37 server.should_exit = True
38 try:
39 await asyncio.wait_for(task, timeout=5)
40 finally:
41 sock.close()
42
43
44 @pytest.mark.asyncio
45 async def test_socketio_wildcard_handler_only_runs_for_unhandled_events() -> None:
46 import socketio
47
48 handled_calls: list[tuple[str, Any]] = []
49 wildcard_calls: list[str] = []
50
51 sio = socketio.AsyncServer(async_mode="asgi", cors_allowed_origins="*")
52
53 @sio.on("handled", namespace="/ns")
54 async def _handled(sid: str, data: Any) -> dict[str, Any]:
55 handled_calls.append((sid, data))
56 return {"path": "handled"}
57
58 @sio.on("*", namespace="/ns")
59 async def _wildcard(event: str, sid: str, data: Any) -> dict[str, Any]:
60 wildcard_calls.append(event)
61 return {"path": "wildcard", "event": event}
62
63 app = socketio.ASGIApp(sio)
64
65 async with _run_asgi_app(app) as base_url:
66 client = socketio.AsyncClient()
67 await client.connect(base_url, namespaces=["/ns"])
68 try:
69 res = await client.call("handled", {"x": 1}, namespace="/ns", timeout=2)
70 assert res == {"path": "handled"}
71 assert wildcard_calls == []
72
73 res2 = await client.call("unhandled_event", {"x": 2}, namespace="/ns", timeout=2)
74 assert res2 == {"path": "wildcard", "event": "unhandled_event"}
75 assert wildcard_calls == ["unhandled_event"]
76 finally:
77 await client.disconnect()
78
79
80 @pytest.mark.asyncio
81 async def test_socketio_handler_return_values_ack_only_when_client_requests_ack() -> None:
82 import socketio
83
84 sio = socketio.AsyncServer(async_mode="asgi", cors_allowed_origins="*")
85 sent_packets: list[Any] = []
86
87 original_send_packet = sio._send_packet
88
89 async def _record_send_packet(eio_sid: str, pkt: Any) -> None:
90 sent_packets.append(pkt)
91 await original_send_packet(eio_sid, pkt)
92
93 sio._send_packet = _record_send_packet # type: ignore[assignment]
94
95 @sio.on("returns_value", namespace="/ns")
96 async def _returns_value(_sid: str, _data: Any) -> dict[str, Any]:
97 return {"ok": True}
98
99 app = socketio.ASGIApp(sio)
100
101 async with _run_asgi_app(app) as base_url:
102 client = socketio.AsyncClient()
103 await client.connect(base_url, namespaces=["/ns"])
104 try:
105 sent_packets.clear()
106 await client.emit("returns_value", {"x": 1}, namespace="/ns")
107 await asyncio.sleep(0.05)
108 ack_packets = [p for p in sent_packets if getattr(p, "packet_type", None) in (3, 6)]
109 assert ack_packets == []
110
111 sent_packets.clear()
112 res = await client.call("returns_value", {"x": 2}, namespace="/ns", timeout=2)
113 assert res == {"ok": True}
114 ack_packets = [p for p in sent_packets if getattr(p, "packet_type", None) in (3, 6)]
115 assert len(ack_packets) >= 1
116 finally:
117 await client.disconnect()