187 lines
6.3 KiB
Python
187 lines
6.3 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""WebSocket 日志推送端点测试
|
||
|
||
测试 /ws/logs/{execution_id} 端点的连接、日志回放、实时推送和断开行为。
|
||
利用 TaskExecutor 已有的 subscribe/broadcast 机制进行验证。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
|
||
import pytest
|
||
from starlette.testclient import TestClient
|
||
from starlette.websockets import WebSocketDisconnect
|
||
|
||
from app.main import app
|
||
from app.services.task_executor import task_executor
|
||
|
||
|
||
@pytest.fixture(autouse=True)
|
||
def _cleanup_executor():
|
||
"""每个测试前后清理 TaskExecutor 内部状态。"""
|
||
yield
|
||
# 清理所有残留的缓冲区和订阅者
|
||
for eid in list(task_executor._log_buffers.keys()):
|
||
task_executor.cleanup(eid)
|
||
task_executor._subscribers.clear()
|
||
task_executor._log_buffers.clear()
|
||
|
||
|
||
class TestWebSocketConnection:
|
||
"""WebSocket 连接/断开基本行为"""
|
||
|
||
def test_connect_and_disconnect(self):
|
||
"""客户端能成功建立和关闭 WebSocket 连接。"""
|
||
client = TestClient(app)
|
||
with client.websocket_connect("/ws/logs/test-exec-001") as ws:
|
||
# 连接成功，直接关闭
|
||
pass # __exit__ 会关闭连接
|
||
|
||
def test_connect_registers_subscriber(self):
|
||
"""连接后 TaskExecutor 应注册订阅者。"""
|
||
client = TestClient(app)
|
||
# 预先初始化缓冲区（模拟有任务在运行）
|
||
task_executor._log_buffers["test-exec-002"] = []
|
||
|
||
with client.websocket_connect("/ws/logs/test-exec-002"):
|
||
# 连接期间应有订阅者
|
||
assert "test-exec-002" in task_executor._subscribers
|
||
assert len(task_executor._subscribers["test-exec-002"]) >= 1
|
||
|
||
|
||
class TestLogReplay:
|
||
"""历史日志回放"""
|
||
|
||
def test_replay_existing_logs(self):
|
||
"""连接时应先收到内存缓冲区中已有的日志行。"""
|
||
eid = "test-exec-replay"
|
||
# 预填充日志缓冲区
|
||
task_executor._log_buffers[eid] = [
|
||
"[stdout] 第一行",
|
||
"[stdout] 第二行",
|
||
"[stderr] 警告信息",
|
||
]
|
||
|
||
client = TestClient(app)
|
||
with client.websocket_connect(f"/ws/logs/{eid}") as ws:
|
||
# 应按顺序收到 3 条历史日志
|
||
msg1 = ws.receive_text()
|
||
msg2 = ws.receive_text()
|
||
msg3 = ws.receive_text()
|
||
|
||
assert msg1 == "[stdout] 第一行"
|
||
assert msg2 == "[stdout] 第二行"
|
||
assert msg3 == "[stderr] 警告信息"
|
||
|
||
def test_no_logs_no_replay(self):
|
||
"""没有历史日志时不应收到回放消息。"""
|
||
eid = "test-exec-empty"
|
||
task_executor._log_buffers[eid] = []
|
||
|
||
client = TestClient(app)
|
||
with client.websocket_connect(f"/ws/logs/{eid}") as ws:
|
||
# 发送结束信号让连接正常关闭
|
||
task_executor._broadcast_end(eid)
|
||
# 不应有回放消息，直接收到的是后续的结束信号处理
|
||
|
||
|
||
class TestBroadcastReceive:
|
||
"""实时日志广播接收"""
|
||
|
||
def test_receive_broadcast_messages(self):
|
||
"""连接后应能收到 TaskExecutor 广播的实时日志。"""
|
||
eid = "test-exec-broadcast"
|
||
task_executor._log_buffers[eid] = []
|
||
|
||
client = TestClient(app)
|
||
with client.websocket_connect(f"/ws/logs/{eid}") as ws:
|
||
# 模拟 TaskExecutor 广播日志
|
||
task_executor._broadcast(eid, "[stdout] 实时日志行1")
|
||
task_executor._broadcast(eid, "[stderr] 实时错误行")
|
||
|
||
msg1 = ws.receive_text()
|
||
msg2 = ws.receive_text()
|
||
|
||
assert msg1 == "[stdout] 实时日志行1"
|
||
assert msg2 == "[stderr] 实时错误行"
|
||
|
||
def test_end_signal_closes_connection(self):
|
||
"""收到 None 结束信号后 WebSocket 应正常关闭。"""
|
||
eid = "test-exec-end"
|
||
task_executor._log_buffers[eid] = []
|
||
|
||
client = TestClient(app)
|
||
with client.websocket_connect(f"/ws/logs/{eid}") as ws:
|
||
# 广播一条日志后发送结束信号
|
||
task_executor._broadcast(eid, "[stdout] 最后一行")
|
||
task_executor._broadcast_end(eid)
|
||
|
||
msg = ws.receive_text()
|
||
assert msg == "[stdout] 最后一行"
|
||
|
||
def test_replay_then_broadcast(self):
|
||
"""先回放历史日志，再接收实时广播。"""
|
||
eid = "test-exec-mixed"
|
||
task_executor._log_buffers[eid] = ["[stdout] 历史行"]
|
||
|
||
client = TestClient(app)
|
||
with client.websocket_connect(f"/ws/logs/{eid}") as ws:
|
||
# 先收到历史回放
|
||
replay = ws.receive_text()
|
||
assert replay == "[stdout] 历史行"
|
||
|
||
# 再收到实时广播
|
||
task_executor._broadcast(eid, "[stdout] 新行")
|
||
task_executor._broadcast_end(eid)
|
||
|
||
live = ws.receive_text()
|
||
assert live == "[stdout] 新行"
|
||
|
||
|
||
class TestLogBroadcasterUnit:
|
||
"""直接测试 TaskExecutor 的 subscribe/unsubscribe/broadcast 方法
|
||
（作为 LogBroadcaster 功能的单元测试）。
|
||
"""
|
||
|
||
def test_subscribe_creates_queue(self):
|
||
eid = "unit-sub"
|
||
q = task_executor.subscribe(eid)
|
||
assert isinstance(q, asyncio.Queue)
|
||
assert eid in task_executor._subscribers
|
||
task_executor.unsubscribe(eid, q)
|
||
|
||
def test_unsubscribe_removes_queue(self):
|
||
eid = "unit-unsub"
|
||
q = task_executor.subscribe(eid)
|
||
task_executor.unsubscribe(eid, q)
|
||
# 最后一个订阅者移除后，key 也应被清理
|
||
assert eid not in task_executor._subscribers
|
||
|
||
def test_broadcast_delivers_to_all_subscribers(self):
|
||
eid = "unit-multi"
|
||
q1 = task_executor.subscribe(eid)
|
||
q2 = task_executor.subscribe(eid)
|
||
|
||
task_executor._broadcast(eid, "测试消息")
|
||
|
||
assert q1.get_nowait() == "测试消息"
|
||
assert q2.get_nowait() == "测试消息"
|
||
|
||
task_executor.unsubscribe(eid, q1)
|
||
task_executor.unsubscribe(eid, q2)
|
||
|
||
def test_broadcast_end_sends_none(self):
|
||
eid = "unit-end"
|
||
q = task_executor.subscribe(eid)
|
||
|
||
task_executor._broadcast_end(eid)
|
||
|
||
assert q.get_nowait() is None
|
||
task_executor.unsubscribe(eid, q)
|
||
|
||
def test_broadcast_no_subscribers_is_safe(self):
|
||
"""没有订阅者时广播不应报错。"""
|
||
task_executor._broadcast("nonexistent", "无人接收")
|
||
task_executor._broadcast_end("nonexistent")
|