Kris
6c95dc006a
完成渠道模块的单元测试目录搭建，新增多个领域模型、端口、中间件、事件、服务以及基础设施层的单元测试文件，同时补充了conftest.py的环境变量配置，完善测试基础环境。
153 lines
4.3 KiB
Python
153 lines
4.3 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import pytest
|
|
from fastapi import Request
|
|
|
|
from yuxi.channel.interfaces.sse.endpoint import SseEndpoint, _SseConnection
|
|
|
|
|
|
class _FakeRequest:
|
|
def __init__(self, disconnected: bool = False):
|
|
self._disconnected = disconnected
|
|
|
|
async def is_disconnected(self) -> bool:
|
|
return self._disconnected
|
|
|
|
|
|
class _FakeMetrics:
|
|
def __init__(self):
|
|
self.sse_connections = 0
|
|
|
|
async def set_sse_connections(self, count: int) -> None:
|
|
self.sse_connections = count
|
|
|
|
|
|
@pytest.fixture
|
|
def endpoint():
|
|
return SseEndpoint(max_connections=10, metrics=_FakeMetrics())
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_sse_endpoint_start_stop(endpoint):
|
|
await endpoint.start()
|
|
assert endpoint._cleanup_task is not None
|
|
await endpoint.stop()
|
|
assert endpoint._cleanup_task is None
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_sse_endpoint_connection_count(endpoint):
|
|
assert endpoint.connection_count == 0
|
|
endpoint._connections["test"] = [_SseConnection(session_id="test", queue=asyncio.Queue())]
|
|
assert endpoint.connection_count == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_sse_endpoint_subscribe(endpoint):
|
|
await endpoint.start()
|
|
request = _FakeRequest()
|
|
response = await endpoint.subscribe("session-1", request)
|
|
assert response is not None
|
|
assert endpoint.connection_count == 1
|
|
await endpoint.stop()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_sse_endpoint_subscribe_limit_reached(endpoint):
|
|
await endpoint.start()
|
|
endpoint._max_connections = 1
|
|
request = _FakeRequest()
|
|
|
|
# First connection
|
|
await endpoint.subscribe("session-1", request)
|
|
assert endpoint.connection_count == 1
|
|
|
|
# Second connection should fail
|
|
from fastapi import HTTPException
|
|
with pytest.raises(HTTPException) as exc_info:
|
|
await endpoint.subscribe("session-2", request)
|
|
assert exc_info.value.status_code == 503
|
|
|
|
await endpoint.stop()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_sse_endpoint_push_event(endpoint):
|
|
await endpoint.start()
|
|
queue = asyncio.Queue()
|
|
conn = _SseConnection(session_id="session-1", queue=queue)
|
|
endpoint._connections["session-1"] = [conn]
|
|
|
|
delivered = await endpoint.push_event("session-1", {"type": "test", "data": "hello"})
|
|
assert delivered is True
|
|
assert not queue.empty()
|
|
|
|
await endpoint.stop()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_sse_endpoint_push_event_no_connections(endpoint):
|
|
await endpoint.start()
|
|
delivered = await endpoint.push_event("nonexistent", {"type": "test"})
|
|
assert delivered is False
|
|
await endpoint.stop()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_sse_endpoint_broadcast_shutdown(endpoint):
|
|
await endpoint.start()
|
|
queue = asyncio.Queue()
|
|
conn = _SseConnection(session_id="session-1", queue=queue)
|
|
endpoint._connections["session-1"] = [conn]
|
|
|
|
await endpoint.broadcast_shutdown()
|
|
assert not queue.empty()
|
|
msg = queue.get_nowait()
|
|
assert msg["type"] == "shutdown"
|
|
|
|
await endpoint.stop()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_sse_endpoint_cleanup_stale(endpoint):
|
|
await endpoint.start()
|
|
import time
|
|
old_conn = _SseConnection(session_id="session-1", queue=asyncio.Queue(), last_active_at=time.monotonic() - 400)
|
|
endpoint._connections["session-1"] = [old_conn]
|
|
|
|
await endpoint._cleanup_stale()
|
|
assert "session-1" not in endpoint._connections
|
|
|
|
await endpoint.stop()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_sse_endpoint_update_connection_count(endpoint):
|
|
metrics = _FakeMetrics()
|
|
ep = SseEndpoint(max_connections=10, metrics=metrics)
|
|
await ep.start()
|
|
|
|
ep._connections["session-1"] = [_SseConnection(session_id="session-1", queue=asyncio.Queue())]
|
|
await ep._update_connection_count()
|
|
assert metrics.sse_connections == 1
|
|
|
|
await ep.stop()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_sse_endpoint_push_event_queue_full(endpoint):
|
|
await endpoint.start()
|
|
queue = asyncio.Queue(maxsize=1)
|
|
conn = _SseConnection(session_id="session-1", queue=queue)
|
|
endpoint._connections["session-1"] = [conn]
|
|
|
|
# Fill the queue
|
|
queue.put_nowait({"type": "first"})
|
|
|
|
# This should not raise, but skip the full queue
|
|
delivered = await endpoint.push_event("session-1", {"type": "second"})
|
|
assert delivered is True # At least one connection exists
|
|
|
|
await endpoint.stop()
|