145 lines
4.8 KiB
Python
145 lines
4.8 KiB
Python
"""统一 run 生命周期服务的无 DB 回归测试。"""
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import unittest
|
|
from contextlib import contextmanager
|
|
from types import SimpleNamespace
|
|
from unittest.mock import MagicMock, patch
|
|
from uuid import uuid4
|
|
|
|
from web import run_lifecycle
|
|
|
|
|
|
class ClaimRunTests(unittest.TestCase):
|
|
def test_claim_atomically_adds_message_and_marks_running(self) -> None:
|
|
tid, uid = uuid4(), uuid4()
|
|
task = SimpleNamespace(
|
|
task_id=tid,
|
|
user_id=uid,
|
|
run_status="idle",
|
|
model_profile="old",
|
|
)
|
|
session = MagicMock()
|
|
session.execute.side_effect = [
|
|
SimpleNamespace(scalar_one_or_none=lambda: task),
|
|
SimpleNamespace(scalar_one=lambda: 6),
|
|
MagicMock(),
|
|
]
|
|
|
|
@contextmanager
|
|
def fake_scope():
|
|
yield session
|
|
|
|
def prepare(_session, locked_task):
|
|
self.assertIs(locked_task, task)
|
|
return {"model_profile": "new"}, {"marker": "ok"}
|
|
|
|
with patch.object(run_lifecycle, "session_scope", fake_scope):
|
|
claim = run_lifecycle.claim_run_with_message(
|
|
tid, uid, "持久化消息", prepare=prepare,
|
|
attachment_refs=[{"path": "a.png", "kind": "image"}],
|
|
)
|
|
|
|
self.assertEqual(claim.metadata, {"marker": "ok"})
|
|
added = session.add.call_args.args[0]
|
|
self.assertEqual(claim.run_id, added.message_id)
|
|
self.assertEqual(added.task_id, tid)
|
|
self.assertEqual(added.idx, 7)
|
|
self.assertEqual(
|
|
added.payload,
|
|
{"role": "user", "content": "持久化消息"},
|
|
)
|
|
self.assertEqual(added.attachment_refs, [{"path": "a.png", "kind": "image"}])
|
|
update_stmt = session.execute.call_args_list[-1].args[0]
|
|
params = update_stmt.compile().params
|
|
self.assertEqual(params["run_status"], "running")
|
|
self.assertIsNone(params["run_error"])
|
|
self.assertEqual(params["model_profile"], "new")
|
|
|
|
def test_busy_task_is_rejected_before_message_insert(self) -> None:
|
|
session = MagicMock()
|
|
session.execute.return_value = SimpleNamespace(
|
|
scalar_one_or_none=lambda: SimpleNamespace(run_status="running")
|
|
)
|
|
|
|
@contextmanager
|
|
def fake_scope():
|
|
yield session
|
|
|
|
with (
|
|
patch.object(run_lifecycle, "session_scope", fake_scope),
|
|
self.assertRaises(run_lifecycle.RunTaskBusy),
|
|
):
|
|
run_lifecycle.claim_run_with_message(uuid4(), uuid4(), "消息")
|
|
session.add.assert_not_called()
|
|
|
|
|
|
class ScheduleRunTests(unittest.IsolatedAsyncioTestCase):
|
|
async def test_schedule_uses_persisted_entry_and_tracks_inflight(self) -> None:
|
|
tid, uid = uuid4(), uuid4()
|
|
app = SimpleNamespace(state=SimpleNamespace(inflight={}))
|
|
worker = MagicMock()
|
|
broker = MagicMock()
|
|
|
|
with (
|
|
patch.object(run_lifecycle, "run_agent_bg", worker),
|
|
patch.object(run_lifecycle, "broker", broker),
|
|
):
|
|
task = run_lifecycle.schedule_claimed_run(
|
|
app,
|
|
tid,
|
|
uid,
|
|
"已落库",
|
|
image_variant="image-x",
|
|
video_variant="video-y",
|
|
scheduled=True,
|
|
)
|
|
self.assertEqual(app.state.inflight[task], tid)
|
|
await task
|
|
await asyncio.sleep(0)
|
|
|
|
worker.assert_called_once_with(
|
|
tid,
|
|
uid,
|
|
"已落库",
|
|
"image-x",
|
|
"video-y",
|
|
True,
|
|
user_message_persisted=True,
|
|
)
|
|
self.assertEqual(app.state.inflight, {})
|
|
broker.start.assert_called_once_with(tid)
|
|
|
|
async def test_schedule_failure_marks_error_and_closes_broker(self) -> None:
|
|
tid, uid = uuid4(), uuid4()
|
|
app = SimpleNamespace(state=SimpleNamespace(inflight={}))
|
|
session = MagicMock()
|
|
broker = MagicMock()
|
|
|
|
@contextmanager
|
|
def fake_scope():
|
|
yield session
|
|
|
|
with (
|
|
patch.object(run_lifecycle, "session_scope", fake_scope),
|
|
patch.object(run_lifecycle, "broker", broker),
|
|
patch.object(
|
|
run_lifecycle.asyncio,
|
|
"create_task",
|
|
side_effect=RuntimeError("loop closed"),
|
|
),
|
|
self.assertRaises(run_lifecycle.RunScheduleError),
|
|
):
|
|
run_lifecycle.schedule_claimed_run(app, tid, uid, "已落库")
|
|
|
|
update_stmt = session.execute.call_args.args[0]
|
|
params = update_stmt.compile().params
|
|
self.assertEqual(params["run_status"], "error")
|
|
self.assertIn("loop closed", params["run_error"])
|
|
broker.close.assert_called_once_with(tid)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|