91 lines
4.2 KiB
Python
91 lines
4.2 KiB
Python
|
"""Native jobs + service lifecycle, with an explicit in-memory repository."""
|
||
|
|
import asyncio
|
||
|
|
import unittest
|
||
|
|
from copy import deepcopy
|
||
|
|
from types import SimpleNamespace
|
||
|
|
|
||
|
|
from services.background_tasks.service import TaskService
|
||
|
|
from services.background_tasks.records import TERMINAL
|
||
|
|
from schemas import ToolUpsert
|
||
|
|
from routes.tools import _payload
|
||
|
|
|
||
|
|
|
||
|
|
class MemoryRepository:
|
||
|
|
def __init__(self):
|
||
|
|
self.records = {}
|
||
|
|
async def create(self, record):
|
||
|
|
old = next((r for r in self.records.values() if r.origin_key == record.origin_key), None)
|
||
|
|
if old:
|
||
|
|
return deepcopy(old)
|
||
|
|
self.records[record.id] = deepcopy(record)
|
||
|
|
return record
|
||
|
|
async def get(self, task_id):
|
||
|
|
return deepcopy(self.records.get(task_id))
|
||
|
|
async def list(self, assistant_id=None, active=False):
|
||
|
|
return [deepcopy(r) for r in self.records.values() if (not assistant_id or r.assistant_id == assistant_id) and (not active or r.status not in TERMINAL)]
|
||
|
|
async def patch(self, task_id, **changes):
|
||
|
|
record = self.records[task_id]
|
||
|
|
if record.status not in TERMINAL:
|
||
|
|
for key, value in changes.items():
|
||
|
|
setattr(record, key, value)
|
||
|
|
record.sequence += 1
|
||
|
|
return deepcopy(record)
|
||
|
|
|
||
|
|
|
||
|
|
class SimulatedExecutor:
|
||
|
|
release = None
|
||
|
|
calls = 0
|
||
|
|
def __init__(self, config, secrets):
|
||
|
|
pass
|
||
|
|
async def prepare(self):
|
||
|
|
return {"session_id": "ses_test", "message_id": "msg_test"}
|
||
|
|
async def run(self, handle, instruction, emit, resume=False):
|
||
|
|
type(self).calls += 1
|
||
|
|
await emit({"status": "running", "summary": "正在执行测试"})
|
||
|
|
await type(self).release.wait()
|
||
|
|
return {"status": "succeeded", "summary": "完成", "result": "真实Job，模拟远端结果"}
|
||
|
|
async def cancel(self, handle):
|
||
|
|
return True
|
||
|
|
|
||
|
|
|
||
|
|
class TaskServiceTest(unittest.IsolatedAsyncioTestCase):
|
||
|
|
async def asyncSetUp(self):
|
||
|
|
SimulatedExecutor.release = asyncio.Event()
|
||
|
|
SimulatedExecutor.calls = 0
|
||
|
|
self.service = TaskService(MemoryRepository(), SimulatedExecutor)
|
||
|
|
await self.service.start()
|
||
|
|
self.tool = SimpleNamespace(type="task", id="tool_task", definition={"config": {"url": "https://remote.test", "directory": "/demo", "timeout_seconds": 30}}, secrets={"password": "never-public"})
|
||
|
|
async def asyncTearDown(self):
|
||
|
|
await self.service.close()
|
||
|
|
async def wait_status(self, task_id, expected):
|
||
|
|
async with asyncio.timeout(3):
|
||
|
|
while (await self.service.get(task_id)).status != expected:
|
||
|
|
await asyncio.sleep(.005)
|
||
|
|
async def submit(self, key="call_1"):
|
||
|
|
return await self.service.submit(assistant_id="asst_1", conversation_id="conv_1", tool=self.tool, instruction="执行测试", invocation_id=key)
|
||
|
|
async def test_idempotence_public_projection_and_completion(self):
|
||
|
|
task = await self.submit()
|
||
|
|
again = await self.submit()
|
||
|
|
self.assertEqual(task["id"], again["id"])
|
||
|
|
await self.wait_status(task["id"], "running")
|
||
|
|
self.assertEqual(SimulatedExecutor.calls, 1)
|
||
|
|
self.assertNotIn("never-public", str(task))
|
||
|
|
SimulatedExecutor.release.set()
|
||
|
|
await self.wait_status(task["id"], "succeeded")
|
||
|
|
with self.assertRaises(ValueError):
|
||
|
|
await self.service.get(task["id"], "another-assistant")
|
||
|
|
async def test_cancel_requires_remote_confirmation(self):
|
||
|
|
task = await self.submit()
|
||
|
|
await self.wait_status(task["id"], "running")
|
||
|
|
await self.service.cancel(task["id"])
|
||
|
|
await self.wait_status(task["id"], "cancelled")
|
||
|
|
async def test_shutdown_does_not_abort_remote_work(self):
|
||
|
|
task = await self.submit()
|
||
|
|
await self.wait_status(task["id"], "running")
|
||
|
|
await self.service.close()
|
||
|
|
self.assertEqual((await self.service.get(task["id"])).status, "running")
|
||
|
|
def test_task_tool_preserves_masked_password(self):
|
||
|
|
body = ToolUpsert(name="后台编程", functionName="coding_task", definition={"type": "task", "config": {"url": "https://remote.test", "directory": "/demo"}}, secrets={"password": "****"})
|
||
|
|
data = _payload(body, {"password": "old-password"})
|
||
|
|
self.assertEqual(data["secrets"]["password"], "old-password")
|