tests/test_orchestrator.py
Rama main · réplica de lectura del repositorio original.
from __future__ import annotations
import sys
import tempfile
import unittest
from dataclasses import dataclass
from pathlib import Path
from typing import Any
ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT / "src"))
from chorus.backends import RUN_STATUS_SUCCEEDED # noqa: E402
from chorus.orchestrator import Orchestrator # noqa: E402
from chorus.persistence import ChorusStore, open_store # noqa: E402
from chorus.tickets import create_ticket, transition_ticket # noqa: E402
@dataclass
class FakeRunResult:
run_id: str
status: str
exit_code: int | None = None
class FakeRunHandle:
def __init__(self, store: ChorusStore, run_id: str) -> None:
self.store = store
self.run_id = run_id
self._result: FakeRunResult | None = None
@property
def is_done(self) -> bool:
return self._result is not None
def finish(self, status: str, exit_code: int | None = 0) -> None:
self._result = FakeRunResult(self.run_id, status, exit_code)
self.store.runs.update(
self.run_id,
status=status,
exit_code=exit_code,
finished_at="2026-06-15T12:00:00Z",
)
def wait(self, timeout: float | None = None) -> FakeRunResult:
if self._result is None:
raise TimeoutError("Fake run has not finished.")
return self._result
class FakeBackend:
def __init__(self, store: ChorusStore) -> None:
self.store = store
self.starts: list[dict[str, Any]] = []
self.handles: list[FakeRunHandle] = []
def start_task(
self,
*,
project_id: str,
prompt: str,
repo_path: str | Path,
task_id: str | None = None,
agent_id: str | None = None,
) -> FakeRunHandle:
run = self.store.runs.create(
project_id=project_id,
task_id=task_id,
agent_id=agent_id,
backend="fake",
command="fake-agent",
status="running",
log_path=str(Path(repo_path) / ".chorus" / "fake.log"),
started_at="2026-06-15T12:00:00Z",
)
handle = FakeRunHandle(self.store, str(run["id"]))
self.handles.append(handle)
self.starts.append(
{
"project_id": project_id,
"prompt": prompt,
"repo_path": str(repo_path),
"task_id": task_id,
"agent_id": agent_id,
"run_id": run["id"],
}
)
return handle
class OrchestratorLoopTests(unittest.TestCase):
def setUp(self) -> None:
self.tempdir = tempfile.TemporaryDirectory(dir=ROOT)
self.root = Path(self.tempdir.name)
self.repo_path = self.root / "repo"
docs_path = self.repo_path / "docs"
docs_path.mkdir(parents=True)
(docs_path / "SPEC.md").write_text("# Spec\nBuild Chorus.\n", encoding="utf-8")
(docs_path / "ARCHITECTURE.md").write_text(
"# Architecture\nHub and spoke.\n",
encoding="utf-8",
)
self.store = open_store(self.root / "chorus.db")
self.project = self.store.projects.create(
name="Chorus",
repo_url="https://github.com/faridani/chorus",
local_path=str(self.repo_path),
spec_path="docs/SPEC.md",
architecture_path="docs/ARCHITECTURE.md",
status="active",
default_branch="main",
integration_branch="chorus/integration",
)
self.dev_role = self.store.roles.create(
project_id=self.project["id"],
name="software dev",
description="Implements tickets",
allowed_actions_json=["edit_files", "run_tests", "commit"],
forbidden_actions_json=["git_push", "merge_main"],
)
self.qa_role = self.store.roles.create(
project_id=self.project["id"],
name="qa",
description="Verifies behavior",
)
self.agent = self.store.agents.create(
project_id=self.project["id"],
role_id=self.dev_role["id"],
name="dev-1",
backend="fake",
status="idle",
)
self.store.guardrails.create(
project_id=self.project["id"],
scope="project",
name="No main writes",
rules_json=["Never touch main."],
)
self.store.guardrails.create(
project_id=self.project["id"],
role_id=self.dev_role["id"],
scope="role",
name="Commit work",
rules_json=["Commit all changes before finishing."],
)
self.backend = FakeBackend(self.store)
def tearDown(self) -> None:
self.store.close()
self.tempdir.cleanup()
def test_run_once_dispatches_highest_priority_ready_ticket_with_available_role(self) -> None:
unavailable_high = create_ticket(
self.store,
project_id=self.project["id"],
title="QA the release",
body="No QA agent is idle yet.",
assigned_role_id=self.qa_role["id"],
status="ready",
priority=100,
)
selected = create_ticket(
self.store,
project_id=self.project["id"],
title="Implement orchestrator task loop",
body="Turn ready tickets into agent tasks.",
assigned_role_id=self.dev_role["id"],
status="ready",
priority=50,
labels=["backend"],
)
lower = create_ticket(
self.store,
project_id=self.project["id"],
title="Polish dashboard copy",
assigned_role_id=self.dev_role["id"],
status="ready",
priority=1,
)
result = Orchestrator(self.store, backends={"fake": self.backend}).run_once(
self.project["id"]
)
self.assertEqual(1, len(result["started"]))
started = result["started"][0]
self.assertEqual(selected["id"], started["ticket_id"])
self.assertEqual("in_progress", self.store.tickets.require(selected["id"])["status"])
self.assertEqual("ready", self.store.tickets.require(unavailable_high["id"])["status"])
self.assertEqual("ready", self.store.tickets.require(lower["id"])["status"])
task = self.store.tasks.require(started["task_id"])
run = self.store.runs.require(started["run_id"])
agent = self.store.agents.require(self.agent["id"])
branch_states = self.store.branch_states.list(
filters={"project_id": self.project["id"], "branch_name": started["branch_name"]},
order_by=None,
)
self.assertEqual("running", task["status"])
self.assertEqual(selected["id"], task["ticket_id"])
self.assertEqual(self.agent["id"], run["agent_id"])
self.assertEqual("running", run["status"])
self.assertEqual("running", agent["status"])
self.assertEqual(started["branch_name"], agent["branch_name"])
self.assertEqual([started["branch_name"]], [branch["branch_name"] for branch in branch_states])
self.assertEqual("allocated", branch_states[0]["status"])
context = task["context_json"]
self.assertEqual("software dev", context["role"]["name"])
self.assertEqual(selected["id"], context["ticket"]["id"])
self.assertEqual(["backend"], context["ticket"]["labels"])
self.assertEqual(2, len(context["guardrails"]))
self.assertIn("# Spec", context["project_context"]["spec"]["content"])
self.assertIn("Hub and spoke.", self.backend.starts[0]["prompt"])
self.assertIn("Implement orchestrator task loop", self.backend.starts[0]["prompt"])
self.assertEqual(str(self.repo_path), self.backend.starts[0]["repo_path"])
def test_task_context_includes_tool_permission_states_from_catalog(self) -> None:
self.store.projects.update(
self.project["id"],
metadata_json={
"available_actions": [
"edit_files",
{"name": "git_push"},
{"id": "open_pr"},
]
},
)
selected = create_ticket(
self.store,
project_id=self.project["id"],
title="Use structured tool permissions",
assigned_role_id=self.dev_role["id"],
status="ready",
priority=1,
)
result = Orchestrator(self.store, backends={"fake": self.backend}).run_once(
self.project["id"]
)
self.assertEqual(1, len(result["started"]))
task = self.store.tasks.require(result["started_task_ids"][0])
self.assertEqual(selected["id"], task["ticket_id"])
permissions = task["context_json"]["role"]["tool_permissions"]
states = {permission["name"]: permission["state"] for permission in permissions}
self.assertEqual("allowed", states["edit_files"])
self.assertEqual("disallowed", states["git_push"])
self.assertEqual("unspecified", states["open_pr"])
self.assertEqual(
["edit_files", "git_push", "open_pr", "run_tests", "commit", "merge_main"],
[permission["name"] for permission in permissions],
)
self.assertIn("Tool permissions: ", self.backend.starts[0]["prompt"])
self.assertIn(
'"name": "open_pr", "state": "unspecified"',
self.backend.starts[0]["prompt"],
)
def test_run_once_marks_catalog_tools_unspecified_without_role_action_lists(self) -> None:
self.store.projects.update(
self.project["id"],
metadata_json={"available_actions": [{"action": "open_pr"}]},
)
self.store.agents.create(
project_id=self.project["id"],
role_id=self.qa_role["id"],
name="qa-1",
backend="fake",
status="idle",
)
ticket = create_ticket(
self.store,
project_id=self.project["id"],
title="Verify legacy role actions",
assigned_role_id=self.qa_role["id"],
status="ready",
priority=1,
)
result = Orchestrator(self.store, backends={"fake": self.backend}).run_once(
self.project["id"]
)
self.assertEqual(1, len(result["started"]))
task = self.store.tasks.require(result["started_task_ids"][0])
self.assertEqual(ticket["id"], task["ticket_id"])
self.assertEqual(
[{"name": "open_pr", "state": "unspecified"}],
task["context_json"]["role"]["tool_permissions"],
)
self.assertIn(
'Tool permissions: [{"name": "open_pr", "state": "unspecified"}]',
self.backend.starts[0]["prompt"],
)
def test_run_once_reconciles_finished_agent_run(self) -> None:
ticket = create_ticket(
self.store,
project_id=self.project["id"],
title="Finish a task",
assigned_role_id=self.dev_role["id"],
status="ready",
priority=1,
)
orchestrator = Orchestrator(self.store, backends={"fake": self.backend})
first_result = orchestrator.run_once(self.project["id"])
run_id = first_result["started_run_ids"][0]
task_id = first_result["started_task_ids"][0]
self.backend.handles[0].finish(RUN_STATUS_SUCCEEDED)
second_result = orchestrator.run_once(self.project["id"])
task = self.store.tasks.require(task_id)
agent = self.store.agents.require(self.agent["id"])
self.assertEqual([run_id], second_result["reconciled_run_ids"])
self.assertEqual("review", self.store.tickets.require(ticket["id"])["status"])
self.assertEqual("review", task["status"])
self.assertEqual(RUN_STATUS_SUCCEEDED, task["result_json"]["run_status"])
self.assertEqual("idle", agent["status"])
self.assertIsNone(agent["branch_name"])
def test_run_once_does_not_duplicate_persisted_active_work_after_restart(self) -> None:
ticket = create_ticket(
self.store,
project_id=self.project["id"],
title="Already running",
assigned_role_id=self.dev_role["id"],
status="ready",
priority=99,
)
task = self.store.tasks.create(
project_id=self.project["id"],
ticket_id=ticket["id"],
role_id=self.dev_role["id"],
assigned_agent_id=self.agent["id"],
title="Already running",
status="running",
branch_name="chorus/agent/software-dev/already-running",
)
self.store.runs.create(
project_id=self.project["id"],
task_id=task["id"],
agent_id=self.agent["id"],
backend="fake",
command="fake-agent",
status="running",
)
restarted_backend = FakeBackend(self.store)
result = Orchestrator(self.store, backends={"fake": restarted_backend}).run_once(
self.project["id"]
)
self.assertEqual([], result["started_run_ids"])
self.assertEqual([], restarted_backend.starts)
self.assertEqual(1, len(self.store.tasks.list(filters={"ticket_id": ticket["id"]})))
self.assertEqual(1, len(self.store.runs.list(filters={"task_id": task["id"]})))
def test_reconciles_terminal_run_persisted_before_restart(self) -> None:
ticket = create_ticket(
self.store,
project_id=self.project["id"],
title="Terminal before restart",
assigned_role_id=self.dev_role["id"],
status="ready",
priority=1,
)
transition_ticket(self.store, ticket["id"], "in_progress")
task = self.store.tasks.create(
project_id=self.project["id"],
ticket_id=ticket["id"],
role_id=self.dev_role["id"],
assigned_agent_id=self.agent["id"],
title="Terminal before restart",
status="running",
branch_name="chorus/agent/software-dev/terminal-before-restart",
)
run = self.store.runs.create(
project_id=self.project["id"],
task_id=task["id"],
agent_id=self.agent["id"],
backend="fake",
command="fake-agent",
status=RUN_STATUS_SUCCEEDED,
finished_at="2026-06-15T12:00:00Z",
)
result = Orchestrator(self.store, backends={"fake": FakeBackend(self.store)}).run_once(
self.project["id"]
)
self.assertEqual([run["id"]], result["reconciled_run_ids"])
self.assertEqual("review", self.store.tickets.require(ticket["id"])["status"])
self.assertEqual("review", self.store.tasks.require(task["id"])["status"])
self.assertEqual("idle", self.store.agents.require(self.agent["id"])["status"])
if __name__ == "__main__":
unittest.main()