agentci/tests/test_storage.py

import sqlite3
from pathlib import Path

import pytest

from agentci.adapters.storage import Storage
from agentci.domain.events import (
    JobStarted,
    PermissionDenied,
    PermissionGranted,
    WorkflowCreated,
)
from agentci.domain.models import CommandEvent, Workflow, WorkflowKind, WorkflowStatus
from agentci.state_machine import StateMachine


@pytest.fixture
async def storage(tmp_path: Path) -> Storage:
    migrations = Path(__file__).parents[1] / "src" / "agentci" / "migrations"
    value = Storage(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,
) -> CommandEvent:
    return CommandEvent(
        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 test_receive_is_idempotent_without_consuming_sequence(storage: Storage) -> None:
    host = StateMachine(storage)
    first = await host.receive(command("delivery-1"))
    duplicate = await host.receive(command("delivery-1"))
    second = await host.receive(command("delivery-2"))

    assert not first.duplicate
    assert duplicate.duplicate
    assert duplicate.state.id == first.state.id
    assert second.state.receive_sequence == first.state.receive_sequence + 1


async def test_received_job_blocks_later_execute_task_for_same_target(
    storage: Storage,
) -> None:
    host = StateMachine(storage)
    first = (await host.receive(command("delivery-1"))).state
    second = (await host.receive(command("delivery-2"))).state
    await host.evolve("grant-2", PermissionGranted(job_id=second.id))

    assert await storage.claim_task("jobs") is None

    await host.evolve("deny-1", PermissionDenied(job_id=first.id))
    task = await storage.claim_task("jobs")
    assert task is not None
    assert task.job_id == second.id


async def test_received_job_does_not_block_a_different_target(storage: Storage) -> None:
    host = StateMachine(storage)
    await host.receive(command("delivery-1", issue=3))
    second = (await host.receive(command("delivery-2", issue=4))).state
    await host.evolve("grant-2", PermissionGranted(job_id=second.id))

    task = await storage.claim_task("jobs")

    assert task is not None
    assert task.job_id == second.id


async def test_claims_multiple_eligible_targets_without_duplicates(storage: Storage) -> None:
    host = StateMachine(storage)
    first = (await host.receive(command("delivery-1", issue=3))).state
    second = (await host.receive(command("delivery-2", issue=4))).state
    await host.evolve("grant-1", PermissionGranted(job_id=first.id))
    await host.evolve("grant-2", PermissionGranted(job_id=second.id))

    claimed = [await storage.claim_task("jobs"), await storage.claim_task("jobs")]

    assert [task.job_id for task in claimed if task is not None] == [first.id, second.id]
    assert await storage.claim_task("jobs") is None


async def test_started_and_finished_timestamps_are_owned_by_store(storage: Storage) -> None:
    host = StateMachine(storage)
    state = (await host.receive(command("delivery-1"))).state
    state = (await host.evolve("grant", PermissionGranted(job_id=state.id))).state
    state = (await host.evolve("start", JobStarted(job_id=state.id))).state
    with sqlite3.connect(storage.database_path) as connection:
        row = connection.execute(
            "SELECT started_at, finished_at FROM jobs WHERE id=?", (state.id,)
        ).fetchone()
    assert row is not None
    assert row[0] is not None
    assert row[1] is None


async def test_workflow_queries_and_completed_protection(storage: Storage, tmp_path: Path) -> None:
    workflow = Workflow(
        id="workflow-1",
        kind=WorkflowKind.PLAN,
        repo_owner="alice",
        repo_name="repo",
        issue_number=3,
        workspace_path=tmp_path / "repo",
        base_sha="abc",
        artifact="# Plan",
        status=WorkflowStatus.COMPLETED,
    )
    await storage.create_workflow(workflow)
    await storage.fail_job_workflow("missing-job")
    loaded = await storage.latest_workflow("alice", "repo", 3, WorkflowKind.PLAN)
    assert loaded is not None
    assert loaded.status is WorkflowStatus.COMPLETED


async def test_workflow_creation_and_job_link_are_atomic(storage: Storage, tmp_path: Path) -> None:
    host = StateMachine(storage)
    state = (await host.receive(command("delivery-1"))).state
    state = (await host.evolve("grant", PermissionGranted(job_id=state.id))).state
    state = (await host.evolve("start", JobStarted(job_id=state.id))).state
    workflow = Workflow(
        id="workflow-atomic",
        kind=WorkflowKind.PLAN,
        repo_owner="alice",
        repo_name="repo",
        issue_number=3,
        workspace_path=tmp_path / "repo",
        base_sha="abc",
    )
    result = await host.evolve(
        "workflow-created",
        WorkflowCreated(job_id=state.id, workflow=workflow, stage="planning"),
    )
    assert result.state.workflow_id == workflow.id
    assert await storage.get_workflow(workflow.id) is not None


async def test_schema_has_receive_sequence_and_no_version(storage: Storage) -> None:
    with sqlite3.connect(storage.database_path) as connection:
        columns = {row[1] for row in connection.execute("PRAGMA table_info(jobs)")}
    assert "receive_sequence" in columns
    assert "version" not in columns