124 lines
4.4 KiB
Python
124 lines
4.4 KiB
Python
import asyncio
|
|
import sqlite3
|
|
from contextlib import closing
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
|
|
from agentci.engine.events import PermissionDenied, PermissionGranted
|
|
from agentci.engine.model import IncomingCommand, QueueName, TaskKind
|
|
from agentci.engine.repository import Repository
|
|
|
|
MIGRATIONS = Path(__file__).parents[1] / "src" / "agentci" / "migrations"
|
|
|
|
|
|
@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, *, issue: int = 3) -> IncomingCommand:
|
|
return IncomingCommand(
|
|
delivery_id=delivery,
|
|
comment_id=int(delivery.rsplit("-", 1)[-1]),
|
|
repo_owner="alice",
|
|
repo_name="repo",
|
|
issue_number=issue,
|
|
pr_number=None,
|
|
requester="alice",
|
|
body="/agent plan",
|
|
)
|
|
|
|
|
|
async def test_concurrent_duplicate_accepts_create_one_job_and_event(
|
|
repository: Repository,
|
|
) -> None:
|
|
results = await asyncio.gather(
|
|
repository.accept(command("delivery-1")),
|
|
repository.accept(command("delivery-1")),
|
|
)
|
|
|
|
assert sorted(result.duplicate for result in results) == [False, True]
|
|
assert len({result.job.id for result in results}) == 1
|
|
with closing(sqlite3.connect(repository.database_path)) as connection, connection:
|
|
job_count = connection.execute("SELECT COUNT(*) FROM jobs").fetchone()
|
|
event_count = connection.execute("SELECT COUNT(*) FROM job_events").fetchone()
|
|
assert job_count == (1,)
|
|
assert event_count == (1,)
|
|
|
|
second = await repository.accept(command("delivery-2"))
|
|
assert second.job.receive_sequence == results[0].job.receive_sequence + 1
|
|
|
|
|
|
async def test_received_job_blocks_later_execute_task_for_same_target(
|
|
repository: Repository,
|
|
) -> None:
|
|
first = (await repository.accept(command("delivery-1"))).job
|
|
second = (await repository.accept(command("delivery-2"))).job
|
|
await repository.apply("grant-2", PermissionGranted(job_id=second.id))
|
|
|
|
assert await repository.claim_task(QueueName.JOBS) is None
|
|
|
|
await repository.apply("deny-1", PermissionDenied(job_id=first.id))
|
|
task = await repository.claim_task(QueueName.JOBS)
|
|
assert task is not None
|
|
assert task.job_id == second.id
|
|
assert task.kind is TaskKind.EXECUTE
|
|
|
|
|
|
async def test_received_job_does_not_block_a_different_target(
|
|
repository: Repository,
|
|
) -> None:
|
|
await repository.accept(command("delivery-1", issue=3))
|
|
second = (await repository.accept(command("delivery-2", issue=4))).job
|
|
await repository.apply("grant-2", PermissionGranted(job_id=second.id))
|
|
|
|
task = await repository.claim_task(QueueName.JOBS)
|
|
assert task is not None
|
|
assert task.job_id == second.id
|
|
|
|
|
|
async def test_claims_multiple_eligible_targets_without_duplicates(
|
|
repository: Repository,
|
|
) -> None:
|
|
first = (await repository.accept(command("delivery-1", issue=3))).job
|
|
second = (await repository.accept(command("delivery-2", issue=4))).job
|
|
await repository.apply("grant-1", PermissionGranted(job_id=first.id))
|
|
await repository.apply("grant-2", PermissionGranted(job_id=second.id))
|
|
|
|
claimed = [
|
|
await repository.claim_task(QueueName.JOBS),
|
|
await repository.claim_task(QueueName.JOBS),
|
|
]
|
|
|
|
assert [task.job_id for task in claimed if task is not None] == [first.id, second.id]
|
|
assert await repository.claim_task(QueueName.JOBS) is None
|
|
|
|
|
|
async def test_concurrent_claims_do_not_duplicate_task(repository: Repository) -> None:
|
|
await repository.accept(command("delivery-1"))
|
|
|
|
claims = await asyncio.gather(
|
|
repository.claim_task(QueueName.CONTROL),
|
|
repository.claim_task(QueueName.CONTROL),
|
|
)
|
|
|
|
claimed = [task for task in claims if task is not None]
|
|
assert len(claimed) == 1
|
|
assert claimed[0].kind is TaskKind.AUTHORIZE
|
|
assert claimed[0].queue is QueueName.CONTROL
|
|
assert await repository.claim_task(QueueName.CONTROL) is None
|
|
|
|
|
|
async def test_duplicate_event_id_cannot_be_reused_for_another_job(
|
|
repository: Repository,
|
|
) -> None:
|
|
first = (await repository.accept(command("delivery-1", issue=3))).job
|
|
second = (await repository.accept(command("delivery-2", issue=4))).job
|
|
await repository.apply("permission", PermissionGranted(job_id=first.id))
|
|
|
|
with pytest.raises(RuntimeError):
|
|
await repository.apply("permission", PermissionGranted(job_id=second.id))
|