313 lines
13 KiB
Python
313 lines
13 KiB
Python
import asyncio
|
|
import os
|
|
import time
|
|
import unittest
|
|
from datetime import datetime, timedelta, timezone
|
|
from tempfile import NamedTemporaryFile
|
|
from types import SimpleNamespace
|
|
from unittest.mock import AsyncMock, patch
|
|
|
|
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
|
|
|
|
from app.models import Base
|
|
from app.services.ai_orchestration_run_service import AIOrchestrationRunService
|
|
from app.services.signal_source_orchestration_service import SignalSourceOrchestrationService
|
|
|
|
|
|
def _utc_now_naive() -> datetime:
|
|
return datetime.now(timezone.utc).replace(tzinfo=None)
|
|
|
|
|
|
class _FastModeSignalService(SignalSourceOrchestrationService):
|
|
session_factory: async_sessionmaker[AsyncSession] | None = None
|
|
|
|
async def _load_project(self, **_: object):
|
|
return SimpleNamespace(id=3, title="项目A", name="项目A", description="desc", owner_id=1)
|
|
|
|
async def _load_signal(self, **_: object):
|
|
return SimpleNamespace(
|
|
id=9,
|
|
title="新增 GitHub 信号源",
|
|
summary="监控 issue",
|
|
signal_type="新需求",
|
|
recommended_action="enter_only",
|
|
content="新增 GitHub 信号源并自动拆任务",
|
|
)
|
|
|
|
async def _plan_onboarding(self, **_: object):
|
|
return (
|
|
{
|
|
"should_onboard": True,
|
|
"summary": "接入 GitHub 仓库信号源。",
|
|
"source": {
|
|
"source_kind": "github",
|
|
"source_name": "核心仓库",
|
|
"source_identifier": "github:openai/codex",
|
|
"description": "监控 issue 与 PR",
|
|
},
|
|
"worker_roles": ["pm_dispatcher", "tech_architect", "delivery_coach"],
|
|
"openclaw_instruction": "创建接入任务。",
|
|
},
|
|
"planned",
|
|
{
|
|
"used_fallback": False,
|
|
"error_code": None,
|
|
"source_field": "message.content",
|
|
"provider": "glm",
|
|
"model": "GLM-4.7",
|
|
},
|
|
)
|
|
|
|
async def _create_or_reuse_registry(self, **_: object):
|
|
return {
|
|
"status": "created",
|
|
"registry_id": 11,
|
|
"source_kind": "github",
|
|
"source_name": "核心仓库",
|
|
"source_identifier": "github:openai/codex",
|
|
"description": "监控 issue 与 PR",
|
|
}
|
|
|
|
async def _resolve_operator_user(self, **_: object):
|
|
return SimpleNamespace(id=1, role="BOSS", name="Boss", username="boss")
|
|
|
|
async def _dispatch_openclaw(self, **_: object):
|
|
return {
|
|
"status": "completed",
|
|
"result": {"trace_id": "openclaw-123"},
|
|
}
|
|
|
|
@classmethod
|
|
async def _execute_workers_and_finalize(cls, **kwargs: object):
|
|
await asyncio.sleep(0.2)
|
|
assert cls.session_factory is not None
|
|
async with cls.session_factory() as db:
|
|
service = cls(db)
|
|
trace_id = str(kwargs["trace_id"])
|
|
summary = str(kwargs["summary"])
|
|
planner_status = str(kwargs["planner_status"])
|
|
planner_diagnostics = kwargs["planner_diagnostics"]
|
|
registry_result = dict(kwargs["registry_result"])
|
|
openclaw_result = dict(kwargs["openclaw_result"])
|
|
worker_role_keys = list(kwargs["worker_role_keys"])
|
|
worker_results = [
|
|
{
|
|
"worker_role_key": role_key,
|
|
"worker_role_name": role_key,
|
|
"status": "completed",
|
|
"summary": f"{role_key} done",
|
|
}
|
|
for role_key in worker_role_keys
|
|
]
|
|
final_result = service._build_orchestration_result(
|
|
trace_id=trace_id,
|
|
response_mode="fast",
|
|
planner_status=planner_status,
|
|
planner_diagnostics=planner_diagnostics,
|
|
summary=summary,
|
|
registry_result=registry_result,
|
|
source_payload=registry_result,
|
|
openclaw_result=openclaw_result,
|
|
worker_results=worker_results,
|
|
pending_workers=[],
|
|
status="completed",
|
|
run_status="completed",
|
|
)
|
|
await AIOrchestrationRunService(db).update_run(
|
|
tenant_id=int(kwargs["tenant_id"]),
|
|
trace_id=trace_id,
|
|
status="completed",
|
|
planner_status=planner_status,
|
|
openclaw_status=str(openclaw_result.get("status") or ""),
|
|
worker_total=len(worker_role_keys),
|
|
worker_completed=len(worker_role_keys),
|
|
worker_failed=0,
|
|
summary=summary,
|
|
planner_diagnostics=planner_diagnostics if isinstance(planner_diagnostics, dict) else {},
|
|
result_payload=final_result,
|
|
completed_at=_utc_now_naive(),
|
|
)
|
|
await db.commit()
|
|
return final_result
|
|
|
|
|
|
class _ConcurrentWorkerSignalService(SignalSourceOrchestrationService):
|
|
@classmethod
|
|
async def _execute_single_worker_role_isolated(cls, **kwargs: object):
|
|
await asyncio.sleep(0.2)
|
|
role = dict(kwargs["role"])
|
|
return {
|
|
"worker_role_key": str(role["role_key"]),
|
|
"worker_role_name": str(role["role_name"]),
|
|
"status": "completed",
|
|
"summary": f"{role['role_key']} done",
|
|
}
|
|
|
|
|
|
class AIOrchestrationRunServiceTest(unittest.IsolatedAsyncioTestCase):
|
|
async def asyncSetUp(self) -> None:
|
|
tmp = NamedTemporaryFile(prefix="pmflow-orchestration-test-", suffix=".db", delete=False)
|
|
tmp.close()
|
|
self.db_path = tmp.name
|
|
self.engine = create_async_engine(f"sqlite+aiosqlite:///{self.db_path}", future=True)
|
|
self.session_factory = async_sessionmaker(self.engine, expire_on_commit=False, class_=AsyncSession)
|
|
async with self.engine.begin() as conn:
|
|
await conn.run_sync(Base.metadata.create_all)
|
|
_FastModeSignalService.session_factory = self.session_factory
|
|
_ConcurrentWorkerSignalService.session_factory = self.session_factory
|
|
|
|
async def asyncTearDown(self) -> None:
|
|
await self.engine.dispose()
|
|
if getattr(self, "db_path", None):
|
|
try:
|
|
os.remove(self.db_path)
|
|
except FileNotFoundError:
|
|
pass
|
|
|
|
async def test_fast_mode_returns_before_background_workers_finish_and_run_is_queryable(self) -> None:
|
|
async with self.session_factory() as db:
|
|
service = _FastModeSignalService(db)
|
|
started = time.perf_counter()
|
|
result = await service.orchestrate_from_signal(
|
|
tenant_id=1,
|
|
project_id=3,
|
|
signal_id=9,
|
|
initiating_text="新增一个 GitHub 信号源，监控核心仓库 issue。",
|
|
requested_via="pmflow",
|
|
operator_user_id=1,
|
|
dry_run=False,
|
|
response_mode="fast",
|
|
)
|
|
elapsed = time.perf_counter() - started
|
|
|
|
self.assertLess(elapsed, 0.15)
|
|
self.assertEqual(result["run_status"], "running")
|
|
self.assertEqual(len(result["pending_workers"]), 3)
|
|
|
|
async with self.session_factory() as db:
|
|
row = await AIOrchestrationRunService(db).get_run(tenant_id=1, trace_id=result["trace_id"])
|
|
self.assertIsNotNone(row)
|
|
self.assertEqual(row.status, "running")
|
|
|
|
await asyncio.sleep(0.3)
|
|
|
|
async with self.session_factory() as db:
|
|
row = await AIOrchestrationRunService(db).get_run(tenant_id=1, trace_id=result["trace_id"])
|
|
snapshot = AIOrchestrationRunService.serialize_run(row)
|
|
|
|
self.assertEqual(snapshot["run_status"], "completed")
|
|
self.assertEqual(snapshot["worker_completed"], 3)
|
|
self.assertEqual(snapshot["worker_pending"], 0)
|
|
|
|
async def test_complete_mode_worker_execution_is_concurrent(self) -> None:
|
|
async with self.session_factory() as db:
|
|
service = _ConcurrentWorkerSignalService(db)
|
|
initial_result = service._build_orchestration_result(
|
|
trace_id="trace-concurrent",
|
|
response_mode="complete",
|
|
planner_status="planned",
|
|
planner_diagnostics={"used_fallback": False},
|
|
summary="接入 GitHub 仓库信号源。",
|
|
registry_result={"registry_id": 11, "status": "created"},
|
|
source_payload={
|
|
"source_kind": "github",
|
|
"source_name": "核心仓库",
|
|
"source_identifier": "github:openai/codex",
|
|
},
|
|
openclaw_result={"status": "completed", "result": {"trace_id": "openclaw-123"}},
|
|
worker_results=[],
|
|
pending_workers=["pm_dispatcher", "tech_architect", "delivery_coach"],
|
|
status="partial",
|
|
run_status="running",
|
|
)
|
|
await AIOrchestrationRunService(db).create_run(
|
|
tenant_id=1,
|
|
trace_id="trace-concurrent",
|
|
run_type="signal_source_onboarding",
|
|
project_id=3,
|
|
signal_id=9,
|
|
response_mode="complete",
|
|
status="running",
|
|
planner_status="planned",
|
|
openclaw_status="completed",
|
|
worker_total=3,
|
|
worker_completed=0,
|
|
worker_failed=0,
|
|
summary="接入 GitHub 仓库信号源。",
|
|
planner_diagnostics={"used_fallback": False},
|
|
result_payload=initial_result,
|
|
created_by_user_id=1,
|
|
)
|
|
await db.commit()
|
|
|
|
fake_roles = {
|
|
"roles": [
|
|
{"role_key": "pm_dispatcher", "role_name": "PM"},
|
|
{"role_key": "tech_architect", "role_name": "Arch"},
|
|
{"role_key": "delivery_coach", "role_name": "Coach"},
|
|
]
|
|
}
|
|
with patch("app.services.signal_source_orchestration_service.prompt_profile_service.get_profiles", AsyncMock(return_value=fake_roles)):
|
|
started = time.perf_counter()
|
|
final_result = await _ConcurrentWorkerSignalService._execute_workers_and_finalize(
|
|
tenant_id=1,
|
|
trace_id="trace-concurrent",
|
|
project_payload={"id": 3, "title": "项目A", "name": "项目A", "description": "desc", "owner_id": 1},
|
|
signal_payload={"id": 9, "title": "新增 GitHub 信号源", "summary": "监控 issue", "signal_type": "新需求", "recommended_action": "enter_only", "content": "content"},
|
|
registry_result={
|
|
"registry_id": 11,
|
|
"status": "created",
|
|
"source_kind": "github",
|
|
"source_name": "核心仓库",
|
|
"source_identifier": "github:openai/codex",
|
|
"description": "监控 issue 与 PR",
|
|
},
|
|
openclaw_result={"status": "completed", "result": {"trace_id": "openclaw-123"}},
|
|
worker_role_keys=["pm_dispatcher", "tech_architect", "delivery_coach"],
|
|
operator_user_payload={"id": 1, "name": "Boss", "username": "boss", "role": "BOSS"},
|
|
planner_status="planned",
|
|
planner_diagnostics={"used_fallback": False},
|
|
summary="接入 GitHub 仓库信号源。",
|
|
signal_id=None,
|
|
)
|
|
elapsed = time.perf_counter() - started
|
|
|
|
self.assertLess(elapsed, 0.45)
|
|
self.assertEqual(final_result["status"], "completed")
|
|
self.assertEqual(len(final_result["worker_executions"]), 3)
|
|
|
|
async def test_mark_stale_runs_as_interrupted(self) -> None:
|
|
async with self.session_factory() as db:
|
|
service = AIOrchestrationRunService(db)
|
|
await service.create_run(
|
|
tenant_id=1,
|
|
trace_id="trace-stale",
|
|
run_type="signal_source_onboarding",
|
|
project_id=3,
|
|
signal_id=9,
|
|
response_mode="fast",
|
|
status="running",
|
|
planner_status="planned",
|
|
openclaw_status="completed",
|
|
worker_total=3,
|
|
worker_completed=1,
|
|
worker_failed=0,
|
|
summary="stale run",
|
|
planner_diagnostics={"used_fallback": False},
|
|
result_payload={"pending_workers": ["pm_dispatcher"], "run_status": "running"},
|
|
created_by_user_id=1,
|
|
started_at=_utc_now_naive() - timedelta(minutes=11),
|
|
)
|
|
await db.commit()
|
|
|
|
updated = await service.mark_stale_runs_interrupted(older_than_minutes=10)
|
|
await db.commit()
|
|
row = await service.get_run(tenant_id=1, trace_id="trace-stale")
|
|
|
|
self.assertEqual(updated, 1)
|
|
self.assertEqual(row.status, "interrupted")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|