mirror of
https://github.com/jxxghp/MoviePilot.git
synced 2026-08-31 21:17:06 +08:00
refactor: isolate agent task execution data boundary (#6433)
This commit is contained in:
InfinityPacer
@@ -0,0 +1,310 @@
|
||||
"""AgentTask 执行服务的取消、终态和调度清理合同。"""
|
||||
|
||||
import asyncio
|
||||
import threading
|
||||
from collections.abc import Callable
|
||||
from uuid import uuid4
|
||||
|
||||
import pytest
|
||||
|
||||
from app.application.agenttask import AgentTaskExecutionService
|
||||
from app.db.adapters.transaction import TransactionalWriteRunner
|
||||
from app.db.oper.agenttask import AgentTaskOper
|
||||
from app.db.session import SessionFactory, async_session_scope
|
||||
from app.db.worker import DatabaseWorker
|
||||
from app.schemas.exception import DatabaseWorkerClosedError
|
||||
|
||||
|
||||
def _add_task(prefix: str, *, trigger_type: str = "cron"):
|
||||
"""创建一条与其他用例隔离的可执行任务。"""
|
||||
user_id = f"{prefix}-{uuid4().hex}"
|
||||
return AgentTaskOper().add(
|
||||
name=f"{prefix} 检查",
|
||||
content="检查资源并报告",
|
||||
trigger_type=trigger_type,
|
||||
cron_expression="0 * * * *" if trigger_type == "cron" else None,
|
||||
run_at="2099-01-01T00:00:00+08:00" if trigger_type == "date" else None,
|
||||
user_id=user_id,
|
||||
username="admin",
|
||||
session_id=f"session-{user_id}",
|
||||
channel=None,
|
||||
source="api",
|
||||
original_chat_id=None,
|
||||
)
|
||||
|
||||
|
||||
def _build_service(
|
||||
worker: DatabaseWorker,
|
||||
repository: Callable[[object], object] | None = None,
|
||||
) -> AgentTaskExecutionService:
|
||||
"""按生产事务和 worker 边界构造独立服务。"""
|
||||
transaction = TransactionalWriteRunner(
|
||||
sync_session=SessionFactory,
|
||||
async_session=async_session_scope,
|
||||
)
|
||||
return AgentTaskExecutionService(
|
||||
repository=repository or (lambda session: AgentTaskOper(session)),
|
||||
async_executor=worker,
|
||||
sync_transaction=transaction.sync,
|
||||
)
|
||||
|
||||
|
||||
async def _wait_for_worker(
|
||||
worker: DatabaseWorker,
|
||||
predicate: Callable[[], bool],
|
||||
) -> None:
|
||||
"""等待 worker 进入目标状态，超时由测试框架明确失败。"""
|
||||
for _ in range(200):
|
||||
if predicate():
|
||||
return
|
||||
await asyncio.sleep(0.005)
|
||||
raise AssertionError(f"数据库 worker 未进入目标状态: {worker.snapshot()}")
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_cancelled_queued_claim_does_not_create_run() -> None:
|
||||
"""认领尚未开始时取消，应撤销排队工作且不得产生运行记录。"""
|
||||
worker = DatabaseWorker(max_workers=1, capacity=2)
|
||||
await worker.start()
|
||||
occupied = threading.Event()
|
||||
release = threading.Event()
|
||||
blocker = asyncio.create_task(worker.run(
|
||||
lambda: (occupied.set(), release.wait())
|
||||
))
|
||||
await asyncio.to_thread(occupied.wait)
|
||||
task = _add_task("queued-cancel")
|
||||
service = _build_service(worker)
|
||||
|
||||
claim = asyncio.create_task(service.claim(task.id))
|
||||
await _wait_for_worker(worker, lambda: worker.snapshot().queued == 1)
|
||||
claim.cancel()
|
||||
await asyncio.sleep(0)
|
||||
release.set()
|
||||
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await claim
|
||||
await blocker
|
||||
current = AgentTaskOper().get(task.id)
|
||||
assert current.last_status == "waiting"
|
||||
assert current.last_run_id is None
|
||||
assert AgentTaskOper().list_runs(task.id) == []
|
||||
await worker.shutdown()
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_cancelled_started_claim_is_compensated_before_return() -> None:
|
||||
"""认领事务已开始时取消，返回前必须把已提交运行收口为失败。"""
|
||||
worker = DatabaseWorker(max_workers=1, capacity=2)
|
||||
await worker.start()
|
||||
started = threading.Event()
|
||||
release = threading.Event()
|
||||
|
||||
class BlockingRepository:
|
||||
"""在认领已写入但事务尚未提交的位置制造取消窗口。"""
|
||||
|
||||
def __init__(self, session: object) -> None:
|
||||
self._repository = AgentTaskOper(session)
|
||||
|
||||
def __getattr__(self, name: str):
|
||||
return getattr(self._repository, name)
|
||||
|
||||
def begin_run(self, *args, **kwargs):
|
||||
run = self._repository.begin_run(*args, **kwargs)
|
||||
started.set()
|
||||
release.wait()
|
||||
return run
|
||||
|
||||
task = _add_task("started-cancel")
|
||||
service = _build_service(worker, BlockingRepository)
|
||||
claim = asyncio.create_task(service.claim(task.id))
|
||||
await asyncio.to_thread(started.wait)
|
||||
claim.cancel()
|
||||
release.set()
|
||||
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await claim
|
||||
current = AgentTaskOper().get(task.id)
|
||||
runs = AgentTaskOper().list_runs(task.id)
|
||||
assert current.last_status == "failed"
|
||||
assert current.last_result == "Agent 定时任务已取消"
|
||||
assert current.run_count == 1
|
||||
assert len(runs) == 1
|
||||
assert runs[0].status == "failed"
|
||||
await worker.shutdown()
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_claim_retries_transient_worker_overload() -> None:
|
||||
"""一次性任务不得因触发瞬间容量已满而永久丢失。"""
|
||||
worker = DatabaseWorker(max_workers=1, capacity=1)
|
||||
await worker.start()
|
||||
task = _add_task("claim-overload", trigger_type="date")
|
||||
service = _build_service(worker)
|
||||
|
||||
occupied = threading.Event()
|
||||
release = threading.Event()
|
||||
blocker = asyncio.create_task(worker.run(
|
||||
lambda: (occupied.set(), release.wait())
|
||||
))
|
||||
await asyncio.to_thread(occupied.wait)
|
||||
claim = asyncio.create_task(service.claim(task.id))
|
||||
await _wait_for_worker(worker, lambda: worker.snapshot().rejected > 0)
|
||||
assert claim.done() is False
|
||||
release.set()
|
||||
|
||||
claimed = await claim
|
||||
await blocker
|
||||
assert claimed.run is not None
|
||||
current = AgentTaskOper().get(task.id)
|
||||
assert current.last_status == "running"
|
||||
assert current.last_run_id == claimed.run.run_id
|
||||
await service.finalize(claimed.run, success=True, result="完成")
|
||||
await worker.shutdown()
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_repeated_finalize_cancellation_waits_for_single_terminal_write() -> None:
|
||||
"""重复取消不得打断已开始的终态事务或重复累计执行次数。"""
|
||||
worker = DatabaseWorker(max_workers=1, capacity=2)
|
||||
await worker.start()
|
||||
started = threading.Event()
|
||||
release = threading.Event()
|
||||
|
||||
class BlockingRepository:
|
||||
"""在运行终态已写入但事务尚未提交的位置制造重复取消窗口。"""
|
||||
|
||||
def __init__(self, session: object) -> None:
|
||||
self._repository = AgentTaskOper(session)
|
||||
|
||||
def __getattr__(self, name: str):
|
||||
return getattr(self._repository, name)
|
||||
|
||||
def finish_run_outcome(self, *args, **kwargs):
|
||||
outcome = self._repository.finish_run_outcome(*args, **kwargs)
|
||||
started.set()
|
||||
release.wait()
|
||||
return outcome
|
||||
|
||||
task = _add_task("finish-cancel")
|
||||
claim_service = _build_service(worker)
|
||||
claimed = await claim_service.claim(task.id)
|
||||
assert claimed.run is not None
|
||||
service = _build_service(worker, BlockingRepository)
|
||||
finalize = asyncio.create_task(service.finalize(
|
||||
claimed.run,
|
||||
success=True,
|
||||
result="完成",
|
||||
))
|
||||
await asyncio.to_thread(started.wait)
|
||||
finalize.cancel()
|
||||
await asyncio.sleep(0)
|
||||
finalize.cancel()
|
||||
release.set()
|
||||
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await finalize
|
||||
current = AgentTaskOper().get(task.id)
|
||||
runs = AgentTaskOper().list_runs(task.id)
|
||||
assert current.last_status == "success"
|
||||
assert current.run_count == 1
|
||||
assert len(runs) == 1
|
||||
assert runs[0].status == "success"
|
||||
await worker.shutdown()
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_cancelled_queued_finalize_waits_for_terminal_write() -> None:
|
||||
"""终态事务仍在队列时取消，返回前也必须完成唯一一次收口。"""
|
||||
worker = DatabaseWorker(max_workers=1, capacity=2)
|
||||
await worker.start()
|
||||
task = _add_task("queued-finish-cancel")
|
||||
service = _build_service(worker)
|
||||
claimed = await service.claim(task.id)
|
||||
assert claimed.run is not None
|
||||
|
||||
occupied = threading.Event()
|
||||
release = threading.Event()
|
||||
blocker = asyncio.create_task(worker.run(
|
||||
lambda: (occupied.set(), release.wait())
|
||||
))
|
||||
await asyncio.to_thread(occupied.wait)
|
||||
finalize = asyncio.create_task(service.finalize(
|
||||
claimed.run,
|
||||
success=True,
|
||||
result="完成",
|
||||
))
|
||||
await _wait_for_worker(worker, lambda: worker.snapshot().queued == 1)
|
||||
finalize.cancel()
|
||||
release.set()
|
||||
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await finalize
|
||||
await blocker
|
||||
current = AgentTaskOper().get(task.id)
|
||||
runs = AgentTaskOper().list_runs(task.id)
|
||||
assert current.last_status == "success"
|
||||
assert current.run_count == 1
|
||||
assert len(runs) == 1
|
||||
assert runs[0].status == "success"
|
||||
await worker.shutdown()
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_finalize_retries_transient_worker_overload() -> None:
|
||||
"""容量暂满时保留终态 owner，取得 admission 后再提交结果。"""
|
||||
worker = DatabaseWorker(max_workers=1, capacity=1)
|
||||
await worker.start()
|
||||
task = _add_task("finish-overload")
|
||||
service = _build_service(worker)
|
||||
claimed = await service.claim(task.id)
|
||||
assert claimed.run is not None
|
||||
|
||||
occupied = threading.Event()
|
||||
release = threading.Event()
|
||||
blocker = asyncio.create_task(worker.run(
|
||||
lambda: (occupied.set(), release.wait())
|
||||
))
|
||||
await asyncio.to_thread(occupied.wait)
|
||||
finalize = asyncio.create_task(service.finalize(
|
||||
claimed.run,
|
||||
success=True,
|
||||
result="完成",
|
||||
))
|
||||
await _wait_for_worker(worker, lambda: worker.snapshot().rejected > 0)
|
||||
assert finalize.done() is False
|
||||
release.set()
|
||||
|
||||
outcome = await finalize
|
||||
await blocker
|
||||
assert outcome.run_finalized is True
|
||||
assert AgentTaskOper().get(task.id).last_status == "success"
|
||||
assert worker.snapshot().queued == 0
|
||||
assert worker.snapshot().running == 0
|
||||
await worker.shutdown()
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_closed_worker_does_not_finalize_or_remove_schedule() -> None:
|
||||
"""持久化不可用时不得伪造终态或清理运行时调度。"""
|
||||
worker = DatabaseWorker(max_workers=1, capacity=2)
|
||||
await worker.start()
|
||||
task = _add_task("closed-finalize", trigger_type="date")
|
||||
service = _build_service(worker)
|
||||
claimed = await service.claim(task.id)
|
||||
assert claimed.run is not None
|
||||
await worker.shutdown()
|
||||
removed: list[tuple[int, int]] = []
|
||||
|
||||
with pytest.raises(DatabaseWorkerClosedError):
|
||||
await service.finalize(
|
||||
claimed.run,
|
||||
success=True,
|
||||
result="完成",
|
||||
scheduler_generation=3,
|
||||
remove_schedule=lambda task_id, generation, _run_id: (
|
||||
removed.append((task_id, generation)) or True
|
||||
),
|
||||
)
|
||||
|
||||
assert removed == []
|
||||
assert AgentTaskOper().get(task.id).last_status == "running"
|
||||
Reference in New Issue
Block a user