393 lines
13 KiB
Python
393 lines
13 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
|
|
from tests.conftest import SQLiteClock
|
|
|
|
|
|
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
|
|
|
|
|
|
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(
|
|
engine_repository: Repository, delivery: str, *, issue: int = 3
|
|
) -> Job:
|
|
job = (await engine_repository.accept(command(delivery, issue=issue))).job
|
|
job = (await engine_repository.apply(f"{delivery}:grant", PermissionGranted(job_id=job.id))).job
|
|
return (await engine_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")
|
|
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(
|
|
engine_repository: Repository,
|
|
workspace: Path,
|
|
*,
|
|
delivery: str,
|
|
workflow_id: str,
|
|
status: WorkflowStatus = WorkflowStatus.ACTIVE,
|
|
) -> tuple[Job, Workflow]:
|
|
job = await create_running_job(engine_repository, delivery)
|
|
workflow = build_workflow(workflow_id, workspace, status=status)
|
|
result = await engine_repository.apply(
|
|
f"{workflow_id}:created",
|
|
WorkflowCreated(job_id=job.id, workflow=workflow, stage="planning"),
|
|
)
|
|
return result.job, workflow
|
|
|
|
|
|
async def test_repository_stamps_start_and_completion_once(
|
|
engine_repository: Repository,
|
|
sqlite_clock: SQLiteClock,
|
|
) -> None:
|
|
job = (await engine_repository.accept(command("delivery-1"))).job
|
|
job = (await engine_repository.apply("grant", PermissionGranted(job_id=job.id))).job
|
|
sqlite_clock.now = "2026-02-01T00:01:00+00:00"
|
|
job = (await engine_repository.apply("start", JobStarted(job_id=job.id))).job
|
|
sqlite_clock.now = "2026-02-01T00:02:00+00:00"
|
|
job = (
|
|
await engine_repository.apply("complete", JobCompleted(job_id=job.id, comment_body="done"))
|
|
).job
|
|
sqlite_clock.now = "2026-02-01T00:03:00+00:00"
|
|
await engine_repository.apply("comment", CommentLinked(job_id=job.id, comment_id=99))
|
|
|
|
with closing(sqlite3.connect(engine_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(
|
|
engine_repository: Repository,
|
|
tmp_path: Path,
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
job = await create_running_job(engine_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 engine_repository.apply(
|
|
"workflow-rollback:created",
|
|
WorkflowCreated(job_id=job.id, workflow=workflow, stage="planning"),
|
|
)
|
|
|
|
persisted_job = await engine_repository.get_job(job.id)
|
|
persisted_workflow = await engine_repository.get_workflow(workflow.id)
|
|
with closing(sqlite3.connect(engine_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(
|
|
engine_repository: Repository,
|
|
) -> None:
|
|
job = (await engine_repository.accept(command("delivery-1"))).job
|
|
applied = await engine_repository.apply("permission", PermissionGranted(job_id=job.id))
|
|
duplicate = await engine_repository.apply("permission", PermissionGranted(job_id=job.id))
|
|
|
|
with closing(sqlite3.connect(engine_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_save_workflow_updates_the_mutable_snapshot_and_timestamp(
|
|
engine_repository: Repository,
|
|
tmp_path: Path,
|
|
sqlite_clock: SQLiteClock,
|
|
) -> None:
|
|
job, workflow = await attach_workflow(
|
|
engine_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,
|
|
)
|
|
sqlite_clock.now = "2026-02-02T00:00:00+00:00"
|
|
|
|
await engine_repository.save_workflow(updated)
|
|
|
|
with closing(sqlite3.connect(engine_repository.database_path)) as connection, connection:
|
|
updated_at = connection.execute(
|
|
"SELECT updated_at FROM workflows WHERE id=?", (workflow.id,)
|
|
).fetchone()
|
|
assert (job.workflow_id, await engine_repository.get_workflow(workflow.id), updated_at) == (
|
|
workflow.id,
|
|
updated,
|
|
("2026-02-02T00:00:00+00:00",),
|
|
)
|
|
|
|
|
|
async def test_saving_unknown_workflow_fails(engine_repository: Repository, tmp_path: Path) -> None:
|
|
workflow = build_workflow("missing", tmp_path)
|
|
|
|
with pytest.raises(KeyError, match="Unknown workflow"):
|
|
await engine_repository.save_workflow(workflow)
|
|
|
|
|
|
async def test_fail_job_workflow_fails_only_active_workflows(
|
|
engine_repository: Repository,
|
|
tmp_path: Path,
|
|
) -> None:
|
|
active_job, active = await attach_workflow(
|
|
engine_repository,
|
|
tmp_path,
|
|
delivery="delivery-1",
|
|
workflow_id="active-workflow",
|
|
)
|
|
completed_job, completed = await attach_workflow(
|
|
engine_repository,
|
|
tmp_path,
|
|
delivery="delivery-2",
|
|
workflow_id="completed-workflow",
|
|
status=WorkflowStatus.COMPLETED,
|
|
)
|
|
|
|
await engine_repository.fail_job_workflow(active_job.id)
|
|
await engine_repository.fail_job_workflow(completed_job.id)
|
|
|
|
assert (
|
|
(await engine_repository.get_workflow(active.id)).status, # type: ignore[union-attr]
|
|
(await engine_repository.get_workflow(completed.id)).status, # type: ignore[union-attr]
|
|
) == (WorkflowStatus.FAILED, WorkflowStatus.COMPLETED)
|
|
|
|
|
|
async def test_latest_workflow_returns_newest_completed_match(
|
|
engine_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(engine_repository.database_path)) as connection, connection:
|
|
_sqlite.insert_workflow(connection, older, "2026-01-01T00:00:00+00:00")
|
|
_sqlite.insert_workflow(connection, newest, "2026-01-02T00:00:00+00:00")
|
|
_sqlite.insert_workflow(connection, active, "2026-01-03T00:00:00+00:00")
|
|
_sqlite.insert_workflow(connection, wrong_issue, "2026-01-04T00:00:00+00:00")
|
|
|
|
assert await engine_repository.latest_workflow("alice", "repo", 3, WorkflowKind.PLAN) == newest
|
|
|
|
|
|
async def test_workflow_for_pr_returns_newest_implementation_match(
|
|
engine_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(engine_repository.database_path)) as connection, connection:
|
|
_sqlite.insert_workflow(connection, older, "2026-01-01T00:00:00+00:00")
|
|
_sqlite.insert_workflow(connection, newest, "2026-01-02T00:00:00+00:00")
|
|
_sqlite.insert_workflow(connection, plan, "2026-01-03T00:00:00+00:00")
|
|
_sqlite.insert_workflow(connection, wrong_repo, "2026-01-04T00:00:00+00:00")
|
|
|
|
assert await engine_repository.workflow_for_pr("alice", "repo", 17) == newest
|
|
|
|
|
|
async def test_implementation_workflows_are_newest_first_and_require_a_pr(
|
|
engine_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(engine_repository.database_path)) as connection, connection:
|
|
_sqlite.insert_workflow(connection, older, "2026-01-01T00:00:00+00:00")
|
|
_sqlite.insert_workflow(connection, newest, "2026-01-02T00:00:00+00:00")
|
|
_sqlite.insert_workflow(connection, no_pr, "2026-01-03T00:00:00+00:00")
|
|
_sqlite.insert_workflow(connection, plan, "2026-01-04T00:00:00+00:00")
|
|
|
|
assert await engine_repository.implementation_workflows("alice", "repo", 3) == [newest, older]
|