tests/test_persistence.py
Rama main · réplica de lectura del repositorio original.
from __future__ import annotations
import sys
import tempfile
import unittest
from pathlib import Path
ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT / "src"))
from chorus.persistence import NotFoundError, open_store # noqa: E402
class PersistenceStoreTests(unittest.TestCase):
def setUp(self) -> None:
self.tempdir = tempfile.TemporaryDirectory(dir=ROOT)
self.database_path = Path(self.tempdir.name) / "chorus.db"
self.store = open_store(self.database_path)
def tearDown(self) -> None:
self.store.close()
self.tempdir.cleanup()
def test_migrations_create_core_tables(self) -> None:
table_rows = self.store.connection.execute(
"""
SELECT name FROM sqlite_master
WHERE type = 'table'
ORDER BY name
"""
).fetchall()
tables = {row["name"] for row in table_rows}
self.assertTrue(
{
"projects",
"agents",
"roles",
"guardrails",
"tickets",
"tasks",
"runs",
"quota_samples",
"branch_states",
"merges",
"notifications",
"changelog_entries",
"schema_migrations",
}.issubset(tables)
)
migration_rows = self.store.connection.execute(
"SELECT id FROM schema_migrations ORDER BY id"
).fetchall()
self.assertEqual(
["001_core", "002_ticket_body"],
[row["id"] for row in migration_rows],
)
def test_project_role_and_guardrail_create_read_update(self) -> None:
project = self.store.projects.create(
name="Chorus",
repo_url="https://github.com/faridani/chorus",
status="active",
metadata_json={"spec": "SPEC.md"},
)
self.assertEqual("chorus/integration", project["integration_branch"])
self.assertEqual({"spec": "SPEC.md"}, project["metadata_json"])
updated_project = self.store.projects.update(
project["id"],
local_path="/worktrees/chorus",
metadata_json={"spec": "SPEC.md", "architecture": "docs/ARCHITECTURE.md"},
)
self.assertEqual("/worktrees/chorus", updated_project["local_path"])
self.assertEqual("docs/ARCHITECTURE.md", updated_project["metadata_json"]["architecture"])
role = self.store.roles.create(
project_id=project["id"],
name="software dev",
description="Implements tickets",
allowed_actions_json=["edit_files", "run_tests", "commit"],
forbidden_actions_json=["git_push", "merge_main"],
)
self.assertEqual(["git_push", "merge_main"], role["forbidden_actions_json"])
guardrail = self.store.guardrails.create(
project_id=project["id"],
role_id=role["id"],
scope="role",
name="No direct main writes",
rules_json=[{"action": "git push", "effect": "deny"}],
)
disabled = self.store.guardrails.update(guardrail["id"], is_enabled=0)
self.assertEqual(0, disabled["is_enabled"])
self.assertEqual("git push", disabled["rules_json"][0]["action"])
def test_ticket_task_agent_run_and_branch_state_resume_flow(self) -> None:
project = self.store.projects.create(name="Chorus", repo_url="https://github.com/faridani/chorus")
role = self.store.roles.create(project_id=project["id"], name="qa")
agent = self.store.agents.create(
project_id=project["id"],
role_id=role["id"],
name="qa-1",
status="idle",
state_json={"last_seen_ticket": None},
)
ticket = self.store.tickets.create(
project_id=project["id"],
assigned_role_id=role["id"],
title="Verify persistence",
status="ready",
labels_json=["persistence"],
)
task = self.store.tasks.create(
project_id=project["id"],
ticket_id=ticket["id"],
role_id=role["id"],
assigned_agent_id=agent["id"],
title="Exercise persistence flows",
branch_name="chorus/ticket/persistence",
context_json={"restart_safe": True},
)
running_task = self.store.tasks.update(
task["id"],
status="running",
attempt_count=1,
state_json={"phase": "tests"},
)
self.assertEqual("tests", running_task["state_json"]["phase"])
run = self.store.runs.create(
project_id=project["id"],
task_id=task["id"],
agent_id=agent["id"],
status="running",
command="codex exec",
state_json={"session": "abc"},
)
self.store.runs.update(run["id"], status="quota_exhausted", resume_after_at="2026-06-15T12:00:00Z")
self.store.branch_states.create(
project_id=project["id"],
branch_name="chorus/ticket/persistence",
kind="agent",
status="clean",
head_sha="abc123",
)
self.store.close()
self.store = open_store(self.database_path)
resumed_run = self.store.runs.require(run["id"])
resumed_tasks = self.store.tasks.list(filters={"project_id": project["id"], "status": "running"})
branch = self.store.branch_states.list(filters={"project_id": project["id"], "branch_name": "chorus/ticket/persistence"})[0]
self.assertEqual("quota_exhausted", resumed_run["status"])
self.assertEqual("2026-06-15T12:00:00Z", resumed_run["resume_after_at"])
self.assertEqual(task["id"], resumed_tasks[0]["id"])
self.assertEqual("abc123", branch["head_sha"])
def test_merge_changelog_notification_and_quota_flows(self) -> None:
project = self.store.projects.create(name="Chorus", repo_url="https://github.com/faridani/chorus")
role = self.store.roles.create(project_id=project["id"], name="developer")
agent = self.store.agents.create(project_id=project["id"], role_id=role["id"], name="dev-1")
task = self.store.tasks.create(project_id=project["id"], role_id=role["id"], title="Implement schema")
run = self.store.runs.create(project_id=project["id"], task_id=task["id"], agent_id=agent["id"], status="succeeded")
sample = self.store.quota_samples.create(
project_id=project["id"],
agent_id=agent["id"],
quota_limit=1000,
quota_used=125,
quota_remaining=875,
raw_json={"source": "codex"},
)
self.assertEqual(875, sample["quota_remaining"])
self.assertEqual({"source": "codex"}, sample["raw_json"])
merge = self.store.merges.create(
project_id=project["id"],
task_id=task["id"],
run_id=run["id"],
merged_by_agent_id=agent["id"],
source_branch="chorus/ticket/schema",
target_branch="chorus/integration",
status="pending",
)
merged = self.store.merges.update(
merge["id"],
status="merged",
target_sha_after="def456",
summary="Defined core persistence schema.",
)
self.assertEqual("merged", merged["status"])
changelog = self.store.changelog_entries.create(
project_id=project["id"],
merge_id=merge["id"],
task_id=task["id"],
agent_id=agent["id"],
title="Core persistence schema",
body="Added SQLite-backed state for the orchestrator.",
metadata_json={"tables": 12},
)
notification = self.store.notifications.create(
project_id=project["id"],
merge_id=merge["id"],
changelog_entry_id=changelog["id"],
channel="email",
recipient="human@example.com",
subject="Chorus merge completed",
payload_json={"template": "merge"},
)
sent = self.store.notifications.update(notification["id"], status="sent", sent_at="2026-06-15T13:00:00Z")
self.assertEqual(12, changelog["metadata_json"]["tables"])
self.assertEqual("sent", sent["status"])
self.assertEqual("merge", sent["payload_json"]["template"])
def test_update_missing_record_raises(self) -> None:
with self.assertRaises(NotFoundError):
self.store.projects.update("missing", status="active")
if __name__ == "__main__":
unittest.main()