mirror of
https://hubproxy.babadafafafafa.cn/https://github.com/jxxghp/MoviePilot.git
synced 2026-09-21 00:23:39 +08:00
268 lines
8.9 KiB
Python
268 lines
8.9 KiB
Python
import asyncio
|
||
import threading
|
||
from unittest.mock import patch
|
||
|
||
import pytest
|
||
|
||
from app.agent.tools.base import (
|
||
MoviePilotTool,
|
||
ToolExecutionTimeoutError,
|
||
_blocking_executors,
|
||
_blocking_futures,
|
||
_blocking_retiring_executors,
|
||
close_blocking_executors,
|
||
reopen_blocking_executors,
|
||
shutdown_blocking_executors,
|
||
)
|
||
from app.agent.tools.manager import MoviePilotToolsManager
|
||
from app.runtime.correlation import correlation_scope, get_correlation_id
|
||
|
||
|
||
class SlowAgentTool(MoviePilotTool):
|
||
"""用于验证工具超时保护的慢工具。"""
|
||
|
||
name: str = "slow_agent_tool"
|
||
description: str = "Test slow tool."
|
||
|
||
async def run(self, **kwargs) -> str:
|
||
"""等待足够久以触发测试中的短超时。"""
|
||
await asyncio.sleep(1)
|
||
return "finished"
|
||
|
||
|
||
class BlockingAgentTool(MoviePilotTool):
|
||
"""用于验证阻塞调用并发名额释放时机的工具。"""
|
||
|
||
name: str = "blocking_agent_tool"
|
||
description: str = "Test blocking tool."
|
||
|
||
async def run(self, **kwargs) -> str:
|
||
"""本测试不会直接调用该方法。"""
|
||
return "unused"
|
||
|
||
|
||
@pytest.fixture(autouse=True)
|
||
def _reset_blocking_executor_runtime():
|
||
"""每个用例前后恢复阻塞池门禁，避免进程级 owner 状态串扰。"""
|
||
assert reopen_blocking_executors() is True
|
||
yield
|
||
assert shutdown_blocking_executors(cancel_futures=True) is True
|
||
assert reopen_blocking_executors() is True
|
||
|
||
|
||
def test_arun_raises_timeout_when_tool_exceeds_limit():
|
||
"""底层工具入口应把超时交给宿主策略记录失败终态。"""
|
||
tool = SlowAgentTool(session_id="session-1", user_id="10001")
|
||
|
||
async def _run_tool():
|
||
with patch("app.agent.tools.base.settings.LLM_TOOL_TIMEOUT", 0.05):
|
||
return await tool._arun()
|
||
|
||
with pytest.raises(ToolExecutionTimeoutError, match="超过 0.05 秒"):
|
||
asyncio.run(_run_tool())
|
||
|
||
|
||
def test_http_tool_manager_uses_same_timeout_guard():
|
||
"""HTTP/MCP 工具入口绕过 _arun 时也应复用工具超时保护。"""
|
||
manager = MoviePilotToolsManager(is_admin=True)
|
||
manager.tools = [SlowAgentTool(session_id="session-1", user_id="10001")]
|
||
|
||
async def _call_tool():
|
||
with patch("app.agent.tools.base.settings.LLM_TOOL_TIMEOUT", 0.05):
|
||
return await manager.call_tool("slow_agent_tool", {})
|
||
|
||
result = asyncio.run(_call_tool())
|
||
|
||
assert "工具 slow_agent_tool 执行超时" in result
|
||
|
||
|
||
def test_run_blocking_keeps_bucket_slot_until_worker_finishes():
|
||
"""被取消的阻塞调用在底层线程结束前不应释放同桶并发名额。"""
|
||
tool = BlockingAgentTool(session_id="session-1", user_id="10001")
|
||
started = asyncio.Event()
|
||
release = threading.Event()
|
||
|
||
def _blocking_call() -> str:
|
||
loop.call_soon_threadsafe(started.set)
|
||
release.wait()
|
||
return "done"
|
||
|
||
async def _run_scenario():
|
||
nonlocal loop
|
||
loop = asyncio.get_running_loop()
|
||
with patch.dict(
|
||
"app.agent.tools.base._blocking_semaphores",
|
||
{"subscribe": asyncio.Semaphore(1)},
|
||
):
|
||
task = asyncio.create_task(tool.run_blocking("subscribe", _blocking_call))
|
||
await started.wait()
|
||
task.cancel()
|
||
with pytest.raises(asyncio.CancelledError):
|
||
await task
|
||
|
||
second_task = asyncio.create_task(
|
||
tool.run_blocking("subscribe", lambda: "second")
|
||
)
|
||
await asyncio.sleep(0.05)
|
||
assert not second_task.done()
|
||
|
||
release.set()
|
||
assert await asyncio.wait_for(second_task, timeout=1) == "second"
|
||
|
||
loop = None
|
||
asyncio.run(_run_scenario())
|
||
|
||
|
||
def test_run_blocking_preserves_each_call_context():
|
||
"""长期复用的工具线程必须读取当前调用，而不是首个调用的上下文。"""
|
||
async def _run_scenario():
|
||
observed = []
|
||
for correlation_id in ("request-one", "request-two"):
|
||
with correlation_scope(correlation_id):
|
||
observed.append(
|
||
await MoviePilotTool.run_blocking("web", get_correlation_id)
|
||
)
|
||
return observed
|
||
|
||
assert asyncio.run(_run_scenario()) == ["request-one", "request-two"]
|
||
|
||
|
||
def test_shutdown_blocking_executors_clears_agent_tool_workers():
|
||
"""测试结束清理应关闭 Agent 工具阻塞线程池，避免全量测试退出时等待 worker。"""
|
||
|
||
async def _create_worker():
|
||
await MoviePilotTool.run_blocking("web", lambda: "done")
|
||
|
||
asyncio.run(_create_worker())
|
||
assert "web" in _blocking_executors
|
||
|
||
shutdown_blocking_executors()
|
||
|
||
assert _blocking_executors == {}
|
||
assert _blocking_futures == {}
|
||
assert _blocking_retiring_executors == set()
|
||
|
||
|
||
def test_shutdown_blocking_executors_cancels_queued_workers_and_is_idempotent():
|
||
"""收尾清理应取消尚未开始的排队任务，并允许重复调用。"""
|
||
shutdown_blocking_executors(wait=False, cancel_futures=True)
|
||
started = [threading.Event(), threading.Event()]
|
||
release = threading.Event()
|
||
queued_ran = threading.Event()
|
||
|
||
def _blocking_call(index: int) -> str:
|
||
started[index].set()
|
||
release.wait()
|
||
return f"done-{index}"
|
||
|
||
async def _run_scenario():
|
||
tasks = [
|
||
asyncio.create_task(MoviePilotTool.run_blocking("web", _blocking_call, index))
|
||
for index in range(2)
|
||
]
|
||
for event in started:
|
||
assert await asyncio.wait_for(asyncio.to_thread(event.wait), timeout=1)
|
||
executor = _blocking_executors["web"]
|
||
queued_future = executor.submit(queued_ran.set)
|
||
shutdown_blocking_executors(wait=False, cancel_futures=True)
|
||
shutdown_blocking_executors(wait=False, cancel_futures=True)
|
||
release.set()
|
||
|
||
assert await asyncio.wait_for(asyncio.gather(*tasks), timeout=1) == ["done-0", "done-1"]
|
||
return queued_future
|
||
|
||
queued_future = asyncio.run(_run_scenario())
|
||
|
||
assert _blocking_executors == {}
|
||
assert _blocking_futures == {}
|
||
assert _blocking_retiring_executors == set()
|
||
assert queued_future.cancelled()
|
||
assert not queued_ran.is_set()
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_close_blocking_executors_retains_owner_until_retry() -> None:
|
||
"""同步调用超时后保留 Future/executor，完成后的重复 close 才成功。"""
|
||
started = threading.Event()
|
||
release = threading.Event()
|
||
|
||
def _blocking_call() -> str:
|
||
"""等待测试释放，稳定制造超过关停预算的运行 Future。"""
|
||
started.set()
|
||
release.wait()
|
||
return "done"
|
||
|
||
task = asyncio.create_task(
|
||
MoviePilotTool.run_blocking("web", _blocking_call)
|
||
)
|
||
assert await asyncio.to_thread(started.wait, 1)
|
||
|
||
try:
|
||
assert await close_blocking_executors(
|
||
timeout_seconds=0.01,
|
||
cancel_futures=True,
|
||
) is False
|
||
assert task.done() is False
|
||
assert _blocking_futures
|
||
assert _blocking_retiring_executors
|
||
|
||
with pytest.raises(RuntimeError, match="正在关闭"):
|
||
await MoviePilotTool.run_blocking("web", lambda: "late")
|
||
|
||
release.set()
|
||
assert await asyncio.wait_for(task, timeout=1) == "done"
|
||
assert await close_blocking_executors(
|
||
timeout_seconds=0.01,
|
||
cancel_futures=True,
|
||
) is True
|
||
assert _blocking_futures == {}
|
||
assert _blocking_retiring_executors == set()
|
||
finally:
|
||
release.set()
|
||
if not task.done():
|
||
await asyncio.wait_for(task, timeout=1)
|
||
|
||
|
||
def test_create_agent_config_uses_llm_max_iterations():
|
||
"""Agent 执行配置应把 LLM_MAX_ITERATIONS 传给 LangGraph recursion_limit。"""
|
||
from app.agent.orchestrator import MoviePilotAgent
|
||
from langchain_core.messages import AIMessage
|
||
|
||
class _FakeGraphState:
|
||
"""提供最小 LangGraph 状态替身。"""
|
||
|
||
values = {"messages": [AIMessage(content="ok")]}
|
||
|
||
class _FakeAgent:
|
||
"""记录 ainvoke 收到的 config。"""
|
||
|
||
def __init__(self) -> None:
|
||
self.config = None
|
||
|
||
async def ainvoke(self, _payload, config=None):
|
||
"""保存运行配置供断言。"""
|
||
self.config = config
|
||
|
||
def get_state(self, _config):
|
||
"""返回最小消息状态。"""
|
||
return _FakeGraphState()
|
||
|
||
async def _execute() -> dict:
|
||
agent = MoviePilotAgent(session_id="session-1", user_id="10001")
|
||
fake_agent = _FakeAgent()
|
||
agent._should_stream = lambda: False
|
||
|
||
async def _create_agent(streaming=False):
|
||
"""返回测试替身 Agent。"""
|
||
return fake_agent
|
||
|
||
agent._create_agent = _create_agent
|
||
agent.stream_handler.stop_streaming = lambda: asyncio.sleep(0, result=(False, ""))
|
||
with patch("app.agent.orchestrator.settings.LLM_MAX_ITERATIONS", 7):
|
||
await agent._execute_agent([])
|
||
return fake_agent.config
|
||
|
||
config = asyncio.run(_execute())
|
||
|
||
assert config["recursion_limit"] == 7
|