197 lines
9.9 KiB
Python
197 lines
9.9 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
|
|
confirms_stop = True
|
|
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 type(self).confirms_stop
|
|
async def snapshot(self, handle):
|
|
return {"status": "running", "summary": "远端仍在执行", "result": ""}
|
|
|
|
|
|
class TaskServiceTest(unittest.IsolatedAsyncioTestCase):
|
|
async def asyncSetUp(self):
|
|
SimulatedExecutor.release = asyncio.Event()
|
|
SimulatedExecutor.calls = 0
|
|
SimulatedExecutor.confirms_stop = True
|
|
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")
|
|
async def test_unconfirmed_remote_stop_is_not_reported_cancelled(self):
|
|
SimulatedExecutor.confirms_stop = False
|
|
task = await self.submit()
|
|
await self.wait_status(task["id"], "running")
|
|
await self.service.cancel(task["id"])
|
|
async with asyncio.timeout(2):
|
|
while self.service.runtime.executor.active_jobs:
|
|
await asyncio.sleep(.005)
|
|
self.assertEqual((await self.service.get(task["id"])).status, "cancel_requested")
|
|
|
|
async def test_queued_task_can_be_cancelled_without_remote_submission(self):
|
|
first = await self.submit("call_first")
|
|
await self.wait_status(first["id"], "running")
|
|
second = await self.submit("call_second")
|
|
await self.service.cancel(second["id"])
|
|
await self.wait_status(second["id"], "cancelled")
|
|
self.assertEqual(SimulatedExecutor.calls, 1)
|
|
|
|
async def test_restart_retries_pending_stop_instead_of_resuming_execution(self):
|
|
SimulatedExecutor.confirms_stop = False
|
|
task = await self.submit()
|
|
await self.wait_status(task["id"], "running")
|
|
await self.service.cancel(task["id"])
|
|
async with asyncio.timeout(2):
|
|
while self.service.runtime.executor.active_jobs:
|
|
await asyncio.sleep(.005)
|
|
await self.service.close()
|
|
SimulatedExecutor.confirms_stop = True
|
|
await self.service.start()
|
|
await self.wait_status(task["id"], "cancelled")
|
|
self.assertEqual(SimulatedExecutor.calls, 1)
|
|
|
|
async def test_unconfirmed_remote_stop_blocks_the_next_submission(self):
|
|
SimulatedExecutor.confirms_stop = False
|
|
first = await self.submit("call_first")
|
|
await self.wait_status(first["id"], "running")
|
|
second = await self.submit("call_second")
|
|
await self.service.cancel(first["id"])
|
|
await self.wait_status(second["id"], "failed")
|
|
self.assertEqual(SimulatedExecutor.calls, 1)
|
|
self.assertFalse((await self.service.get(second["id"])).remote_handle)
|
|
|
|
async def test_restart_observes_existing_remote_handle(self):
|
|
task = await self.submit()
|
|
await self.wait_status(task["id"], "running")
|
|
handle = (await self.service.get(task["id"])).remote_handle
|
|
await self.service.close()
|
|
await self.service.start()
|
|
async with asyncio.timeout(2):
|
|
while SimulatedExecutor.calls < 2:
|
|
await asyncio.sleep(.005)
|
|
self.assertEqual((await self.service.get(task["id"])).remote_handle, handle)
|
|
SimulatedExecutor.release.set()
|
|
await self.wait_status(task["id"], "succeeded")
|
|
|
|
async def test_reconcile_preserves_a_pending_question(self):
|
|
task = await self.submit()
|
|
await self.wait_status(task["id"], "running")
|
|
question = {"kind": "question", "id": "que_1", "questions": [{"question": "使用哪个分支？"}]}
|
|
await self.service.repository.patch(task["id"], status="waiting_input", input_request=question)
|
|
result = await self.service.reconcile(task["id"])
|
|
self.assertEqual(result["status"], "waiting_input")
|
|
self.assertEqual(result["inputRequest"], question)
|
|
|
|
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")
|
|
|
|
|
|
class TaskApiTest(unittest.IsolatedAsyncioTestCase):
|
|
async def test_admin_auth_and_real_native_job_through_http_api(self):
|
|
import httpx
|
|
from fastapi import FastAPI
|
|
from unittest.mock import AsyncMock, patch
|
|
from routes import background_tasks as routes
|
|
from services.auth import require_admin
|
|
from db.session import get_session
|
|
SimulatedExecutor.release = asyncio.Event()
|
|
service = TaskService(MemoryRepository(), SimulatedExecutor)
|
|
await service.start()
|
|
app = FastAPI()
|
|
app.include_router(routes.router)
|
|
async def database():
|
|
yield None
|
|
app.dependency_overrides[get_session] = database
|
|
cfg = SimpleNamespace(type="prompt", runtimeMode="pipeline", tool_bindings=[SimpleNamespace(tool_id="tool_task")], tools=[SimpleNamespace(type="task", id="tool_task", definition={"config": {"url": "https://remote.test", "directory": "/demo", "timeout_seconds": 30}}, secrets={})])
|
|
try:
|
|
async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="https://platform.test") as client:
|
|
unauthenticated = await client.get("/api/background-tasks?assistant_id=asst_1")
|
|
self.assertEqual(unauthenticated.status_code, 401)
|
|
app.dependency_overrides[require_admin] = lambda: {"username": "test"}
|
|
with patch.object(routes, "task_service", service), patch.object(routes, "resolve_runtime_config", AsyncMock(return_value=cfg)):
|
|
response = await client.post("/api/background-tasks", json={"assistantId": "asst_1", "toolId": "tool_task", "instruction": "测试请求", "requestId": "request_1"})
|
|
self.assertEqual(response.status_code, 200)
|
|
task_id = response.json()["id"]
|
|
self.assertEqual(response.json()["assistantId"], "asst_1")
|
|
SimulatedExecutor.release.set()
|
|
async with asyncio.timeout(2):
|
|
while (await service.get(task_id)).status != "succeeded":
|
|
await asyncio.sleep(.005)
|
|
task = await client.get(f"/api/background-tasks/{task_id}")
|
|
self.assertEqual(task.json()["status"], "succeeded")
|
|
self.assertNotIn("secrets", task.json())
|
|
finally:
|
|
await service.close()
|