Files
agentci/tests/test_repository.py
T
2026-07-22 23:10:23 +02:00

463 lines
15 KiB
Python

import sqlite3
from contextlib import closing
from dataclasses import replace
from pathlib import Path
from typing import cast
import pytest
from agentci.engine import _sqlite
from agentci.engine.events import (
CommentLinked,
JobCompleted,
JobStarted,
PermissionGranted,
WorkflowCreated,
)
from agentci.engine.model import (
IncomingCommand,
Job,
QueueName,
TaskKind,
TaskRequest,
Workflow,
WorkflowKind,
WorkflowStatus,
)
from agentci.engine.reducer import Transition
from agentci.engine.repository import Repository
MIGRATIONS = Path(__file__).parents[1] / "src" / "agentci" / "migrations"
class TrackingConnection:
def __init__(self) -> None:
self.closed = False
def __enter__(self) -> "TrackingConnection":
return self
def __exit__(self, *_args: object) -> None:
return None
def close(self) -> None:
self.closed = True
class FailingSetupConnection:
def __init__(self) -> None:
self.closed = False
self.row_factory: object | None = None
def execute(self, _statement: str) -> None:
raise sqlite3.OperationalError("pragma failed")
def close(self) -> None:
self.closed = True
@pytest.fixture
async def repository(tmp_path: Path) -> Repository:
value = Repository(tmp_path / "state.sqlite3", MIGRATIONS)
await value.initialize()
return value
def command(
delivery: str,
body: str = "/agent plan",
*,
issue: int = 3,
pr: int | None = None,
) -> IncomingCommand:
return IncomingCommand(
delivery_id=delivery,
comment_id=int(delivery.rsplit("-", 1)[-1]),
repo_owner="alice",
repo_name="repo",
issue_number=issue,
pr_number=pr,
requester="alice",
body=body,
)
async def create_running_job(repository: Repository, delivery: str, *, issue: int = 3) -> Job:
job = (await repository.accept(command(delivery, issue=issue))).job
job = (await repository.apply(f"{delivery}:grant", PermissionGranted(job_id=job.id))).job
return (await repository.apply(f"{delivery}:start", JobStarted(job_id=job.id))).job
async def test_repository_closes_connections_after_success_and_failure(
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
) -> None:
connections = [TrackingConnection(), TrackingConnection()]
monkeypatch.setattr(
_sqlite,
"connect",
lambda _path: cast(sqlite3.Connection, connections.pop(0)),
)
repository = Repository(tmp_path / "state.sqlite3", MIGRATIONS)
successful = cast(TrackingConnection, await repository._run(lambda connection: connection))
def fail(_connection: sqlite3.Connection) -> None:
raise RuntimeError("operation failed")
failing = connections[0]
with pytest.raises(RuntimeError, match="operation failed"):
await repository._run(fail)
assert successful.closed
assert failing.closed
def test_sqlite_setup_closes_connection_after_pragma_failure(
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
) -> None:
connection = FailingSetupConnection()
monkeypatch.setattr(
_sqlite.sqlite3,
"connect",
lambda *_args, **_kwargs: cast(sqlite3.Connection, connection),
)
with pytest.raises(sqlite3.OperationalError, match="pragma failed"):
_sqlite.connect(tmp_path / "state.sqlite3")
assert connection.closed
def build_workflow(
workflow_id: str,
workspace: Path,
*,
kind: WorkflowKind = WorkflowKind.PLAN,
owner: str = "alice",
repo: str = "repo",
issue: int = 3,
pr: int | None = None,
status: WorkflowStatus = WorkflowStatus.ACTIVE,
) -> Workflow:
return Workflow(
id=workflow_id,
kind=kind,
repo_owner=owner,
repo_name=repo,
issue_number=issue,
pr_number=pr,
workspace_path=workspace / workflow_id,
base_sha=f"base-{workflow_id}",
status=status,
)
async def attach_workflow(
repository: Repository,
workspace: Path,
*,
delivery: str,
workflow_id: str,
status: WorkflowStatus = WorkflowStatus.ACTIVE,
) -> tuple[Job, Workflow]:
job = await create_running_job(repository, delivery)
workflow = build_workflow(workflow_id, workspace, status=status)
result = await repository.apply(
f"{workflow_id}:created",
WorkflowCreated(job_id=job.id, workflow=workflow, stage="planning"),
)
return result.job, workflow
def insert_workflow(
connection: sqlite3.Connection,
workflow: Workflow,
created_at: str,
) -> None:
connection.execute(
"""INSERT INTO workflows(
id, kind, repo_owner, repo_name, issue_number, pr_number, base_sha, branch,
workspace_path, primary_session_id, reviewer_session_id, artifact, review_json,
status, runtime, created_at, updated_at
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""",
(
workflow.id,
workflow.kind.value,
workflow.repo_owner,
workflow.repo_name,
workflow.issue_number,
workflow.pr_number,
workflow.base_sha,
workflow.branch,
str(workflow.workspace_path),
workflow.primary_session_id,
workflow.reviewer_session_id,
workflow.artifact,
workflow.review_json,
workflow.status.value,
workflow.runtime,
created_at,
created_at,
),
)
async def test_receive_is_idempotent_without_consuming_sequence(
repository: Repository,
) -> None:
first = await repository.accept(command("delivery-1"))
duplicate = await repository.accept(command("delivery-1"))
second = await repository.accept(command("delivery-2"))
assert (
first.duplicate,
duplicate.duplicate,
duplicate.job.id,
second.job.receive_sequence,
) == (False, True, first.job.id, first.job.receive_sequence + 1)
async def test_repository_stamps_start_and_completion_once(
repository: Repository,
monkeypatch: pytest.MonkeyPatch,
) -> None:
clock = {"now": "2026-02-01T00:00:00+00:00"}
monkeypatch.setattr(_sqlite, "now", lambda: clock["now"])
job = (await repository.accept(command("delivery-1"))).job
job = (await repository.apply("grant", PermissionGranted(job_id=job.id))).job
clock["now"] = "2026-02-01T00:01:00+00:00"
job = (await repository.apply("start", JobStarted(job_id=job.id))).job
clock["now"] = "2026-02-01T00:02:00+00:00"
job = (await repository.apply("complete", JobCompleted(job_id=job.id, comment_body="done"))).job
clock["now"] = "2026-02-01T00:03:00+00:00"
await repository.apply("comment", CommentLinked(job_id=job.id, comment_id=99))
with closing(sqlite3.connect(repository.database_path)) as connection, connection:
timestamps = connection.execute(
"SELECT started_at, finished_at FROM jobs WHERE id=?", (job.id,)
).fetchone()
assert timestamps == (
"2026-02-01T00:01:00+00:00",
"2026-02-01T00:02:00+00:00",
)
async def test_apply_rolls_back_event_job_workflow_and_task_after_mid_apply_failure(
repository: Repository,
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
) -> None:
job = await create_running_job(repository, "delivery-1")
workflow = build_workflow("workflow-rollback", tmp_path)
original_insert_tasks = _sqlite.insert_tasks
def fail_after_task_write(
connection: sqlite3.Connection,
event_id: str,
transition: Transition,
timestamp: str,
) -> None:
forced_task = TaskRequest(TaskKind.RECONCILE_COMMENT, QueueName.CONTROL)
original_insert_tasks(
connection,
event_id,
replace(transition, tasks=(forced_task,)),
timestamp,
)
raise RuntimeError("injected task persistence failure")
monkeypatch.setattr(_sqlite, "insert_tasks", fail_after_task_write)
with pytest.raises(RuntimeError, match="injected task persistence failure"):
await repository.apply(
"workflow-rollback:created",
WorkflowCreated(job_id=job.id, workflow=workflow, stage="planning"),
)
persisted_job = await repository.get_job(job.id)
persisted_workflow = await repository.get_workflow(workflow.id)
with closing(sqlite3.connect(repository.database_path)) as connection, connection:
event_count = connection.execute(
"SELECT COUNT(*) FROM job_events WHERE event_id=?",
("workflow-rollback:created",),
).fetchone()
task_count = connection.execute(
"SELECT COUNT(*) FROM listener_tasks WHERE source_event_id=?",
("workflow-rollback:created",),
).fetchone()
assert (persisted_job, persisted_workflow, event_count, task_count) == (
job,
None,
(0,),
(0,),
)
async def test_duplicate_event_application_does_not_duplicate_tasks(
repository: Repository,
) -> None:
job = (await repository.accept(command("delivery-1"))).job
applied = await repository.apply("permission", PermissionGranted(job_id=job.id))
duplicate = await repository.apply("permission", PermissionGranted(job_id=job.id))
with closing(sqlite3.connect(repository.database_path)) as connection, connection:
tasks = connection.execute(
"""SELECT ordinal, listener, queue FROM listener_tasks
WHERE source_event_id=? ORDER BY ordinal""",
("permission",),
).fetchall()
assert (duplicate.duplicate, duplicate.job, tasks) == (
True,
applied.job,
[(0, "execute", "jobs"), (1, "reconcile_comment", "control")],
)
async def test_workflow_creation_and_job_link_commit_together(
repository: Repository,
tmp_path: Path,
) -> None:
job, workflow = await attach_workflow(
repository,
tmp_path,
delivery="delivery-1",
workflow_id="workflow-atomic",
)
assert (job.workflow_id, await repository.get_workflow(workflow.id)) == (
workflow.id,
workflow,
)
async def test_save_workflow_updates_the_mutable_snapshot_and_timestamp(
repository: Repository,
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
) -> None:
_, workflow = await attach_workflow(
repository,
tmp_path,
delivery="delivery-1",
workflow_id="workflow-update",
)
updated = replace(
workflow,
pr_number=17,
branch="agent/updated",
primary_session_id="primary",
reviewer_session_id="reviewer",
artifact="# Updated plan",
review_json='{"verdict":"approved"}',
status=WorkflowStatus.COMPLETED,
)
monkeypatch.setattr(_sqlite, "now", lambda: "2026-02-02T00:00:00+00:00")
await repository.save_workflow(updated)
with closing(sqlite3.connect(repository.database_path)) as connection, connection:
updated_at = connection.execute(
"SELECT updated_at FROM workflows WHERE id=?", (workflow.id,)
).fetchone()
assert (await repository.get_workflow(workflow.id), updated_at) == (
updated,
("2026-02-02T00:00:00+00:00",),
)
async def test_saving_unknown_workflow_fails(repository: Repository, tmp_path: Path) -> None:
workflow = build_workflow("missing", tmp_path)
with pytest.raises(KeyError, match="Unknown workflow"):
await repository.save_workflow(workflow)
async def test_fail_job_workflow_fails_only_active_workflows(
repository: Repository,
tmp_path: Path,
) -> None:
active_job, active = await attach_workflow(
repository,
tmp_path,
delivery="delivery-1",
workflow_id="active-workflow",
)
completed_job, completed = await attach_workflow(
repository,
tmp_path,
delivery="delivery-2",
workflow_id="completed-workflow",
status=WorkflowStatus.COMPLETED,
)
await repository.fail_job_workflow(active_job.id)
await repository.fail_job_workflow(completed_job.id)
assert (
(await repository.get_workflow(active.id)).status, # type: ignore[union-attr]
(await repository.get_workflow(completed.id)).status, # type: ignore[union-attr]
) == (WorkflowStatus.FAILED, WorkflowStatus.COMPLETED)
async def test_latest_workflow_returns_newest_completed_match(
repository: Repository,
tmp_path: Path,
) -> None:
older = build_workflow("plan-older", tmp_path, status=WorkflowStatus.COMPLETED)
newest = build_workflow("plan-newest", tmp_path, status=WorkflowStatus.COMPLETED)
active = build_workflow("plan-active", tmp_path)
wrong_issue = build_workflow(
"plan-wrong-issue", tmp_path, issue=4, status=WorkflowStatus.COMPLETED
)
with closing(sqlite3.connect(repository.database_path)) as connection, connection:
insert_workflow(connection, older, "2026-01-01T00:00:00+00:00")
insert_workflow(connection, newest, "2026-01-02T00:00:00+00:00")
insert_workflow(connection, active, "2026-01-03T00:00:00+00:00")
insert_workflow(connection, wrong_issue, "2026-01-04T00:00:00+00:00")
assert await repository.latest_workflow("alice", "repo", 3, WorkflowKind.PLAN) == newest
async def test_workflow_for_pr_returns_newest_implementation_match(
repository: Repository,
tmp_path: Path,
) -> None:
older = build_workflow("implementation-older", tmp_path, kind=WorkflowKind.IMPLEMENT, pr=17)
newest = build_workflow("implementation-newest", tmp_path, kind=WorkflowKind.IMPLEMENT, pr=17)
plan = build_workflow("plan-same-pr", tmp_path, pr=17)
wrong_repo = build_workflow(
"implementation-wrong-repo",
tmp_path,
kind=WorkflowKind.IMPLEMENT,
repo="another",
pr=17,
)
with closing(sqlite3.connect(repository.database_path)) as connection, connection:
insert_workflow(connection, older, "2026-01-01T00:00:00+00:00")
insert_workflow(connection, newest, "2026-01-02T00:00:00+00:00")
insert_workflow(connection, plan, "2026-01-03T00:00:00+00:00")
insert_workflow(connection, wrong_repo, "2026-01-04T00:00:00+00:00")
assert await repository.workflow_for_pr("alice", "repo", 17) == newest
async def test_implementation_workflows_are_newest_first_and_require_a_pr(
repository: Repository,
tmp_path: Path,
) -> None:
older = build_workflow("implementation-older", tmp_path, kind=WorkflowKind.IMPLEMENT, pr=17)
newest = build_workflow("implementation-newest", tmp_path, kind=WorkflowKind.IMPLEMENT, pr=18)
no_pr = build_workflow("implementation-no-pr", tmp_path, kind=WorkflowKind.IMPLEMENT)
plan = build_workflow("plan-with-pr", tmp_path, pr=19)
with closing(sqlite3.connect(repository.database_path)) as connection, connection:
insert_workflow(connection, older, "2026-01-01T00:00:00+00:00")
insert_workflow(connection, newest, "2026-01-02T00:00:00+00:00")
insert_workflow(connection, no_pr, "2026-01-03T00:00:00+00:00")
insert_workflow(connection, plan, "2026-01-04T00:00:00+00:00")
assert await repository.implementation_workflows("alice", "repo", 3) == [newest, older]