mirror of
https://github.com/jxxghp/MoviePilot.git
synced 2026-08-29 03:56:43 +08:00
102 lines
3.2 KiB
Python
102 lines
3.2 KiB
Python
"""Agent 渠道流式输出的 start/stop owner 生命周期测试。"""
|
|
|
|
import asyncio
|
|
|
|
import pytest
|
|
|
|
from app.agent.callback import StreamingHandler
|
|
from app.schemas.types import NotificationChannel
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_repeated_streaming_start_retains_previous_flush_owner() -> None:
|
|
"""重复启动必须先等待旧 flush owner 结束，再发布新一轮上下文。"""
|
|
handler = StreamingHandler()
|
|
handler._can_stream = lambda: True
|
|
handler._source = "old-source"
|
|
handler._streaming_enabled = True
|
|
old_started = asyncio.Event()
|
|
release_old = asyncio.Event()
|
|
observed_sources: list[str | None] = []
|
|
|
|
async def old_flush_owner() -> None:
|
|
"""在退出前读取 handler 上下文，暴露被新一轮提前覆盖的风险。"""
|
|
old_started.set()
|
|
await release_old.wait()
|
|
observed_sources.append(handler._source)
|
|
|
|
old_task = asyncio.create_task(old_flush_owner(), name="old.flush")
|
|
handler._flush_task = old_task
|
|
await old_started.wait()
|
|
|
|
restart = asyncio.create_task(
|
|
handler.start_streaming(
|
|
channel=NotificationChannel.Feishu.value,
|
|
source="new-source",
|
|
),
|
|
name="agent.streaming.restart",
|
|
)
|
|
await asyncio.sleep(0)
|
|
|
|
assert restart.done() is False
|
|
assert handler._source == "old-source"
|
|
assert handler._flush_task is old_task
|
|
|
|
release_old.set()
|
|
await asyncio.wait_for(restart, timeout=1)
|
|
new_task = handler._flush_task
|
|
|
|
assert old_task.done()
|
|
assert observed_sources == ["old-source"]
|
|
assert new_task is not None
|
|
assert new_task is not old_task
|
|
assert handler._source == "new-source"
|
|
|
|
await asyncio.wait_for(handler.stop_streaming(), timeout=1)
|
|
assert new_task.done()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_streaming_start_waits_for_concurrent_stop_final_flush() -> None:
|
|
"""停止阶段的最终刷新未完成时，新启动不得改写消息上下文。"""
|
|
handler = StreamingHandler()
|
|
handler._can_stream = lambda: True
|
|
handler._source = "stopping-source"
|
|
handler._streaming_enabled = True
|
|
flush_started = asyncio.Event()
|
|
release_flush = asyncio.Event()
|
|
|
|
async def blocking_flush() -> None:
|
|
"""把最终刷新停留在生命周期锁内，验证新启动等待。"""
|
|
flush_started.set()
|
|
await release_flush.wait()
|
|
|
|
handler._flush = blocking_flush
|
|
stop_task = asyncio.create_task(
|
|
handler.stop_streaming(),
|
|
name="agent.streaming.stop",
|
|
)
|
|
await flush_started.wait()
|
|
start_task = asyncio.create_task(
|
|
handler.start_streaming(
|
|
channel=NotificationChannel.Feishu.value,
|
|
source="next-source",
|
|
),
|
|
name="agent.streaming.next",
|
|
)
|
|
await asyncio.sleep(0)
|
|
|
|
assert start_task.done() is False
|
|
assert handler._source == "stopping-source"
|
|
|
|
release_flush.set()
|
|
await asyncio.wait_for(stop_task, timeout=1)
|
|
await asyncio.wait_for(start_task, timeout=1)
|
|
next_flush_task = handler._flush_task
|
|
|
|
assert handler._source == "next-source"
|
|
assert next_flush_task is not None
|
|
|
|
await asyncio.wait_for(handler.stop_streaming(), timeout=1)
|
|
assert next_flush_task.done()
|