129 lines
4.8 KiB
Python
129 lines
4.8 KiB
Python
import asyncio
|
|
import sqlite3
|
|
import threading
|
|
from collections.abc import Callable, Coroutine
|
|
from contextlib import closing
|
|
from typing import Any
|
|
|
|
import pytest
|
|
|
|
from agentci.engine.events import PermissionDenied, PermissionGranted
|
|
from agentci.engine.model import IncomingCommand, QueueName, TaskKind
|
|
from agentci.engine.repository import Repository
|
|
|
|
|
|
async def run_concurrently[T](
|
|
*factories: Callable[[], Coroutine[Any, Any, T]],
|
|
) -> list[T]:
|
|
barrier = threading.Barrier(len(factories))
|
|
|
|
def run(factory: Callable[[], Coroutine[Any, Any, T]]) -> T:
|
|
barrier.wait()
|
|
return asyncio.run(factory())
|
|
|
|
return list(await asyncio.gather(*(asyncio.to_thread(run, factory) for factory in factories)))
|
|
|
|
|
|
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(
|
|
engine_repository: Repository,
|
|
) -> None:
|
|
results = await run_concurrently(
|
|
lambda: engine_repository.accept(command("delivery-1")),
|
|
lambda: engine_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(engine_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 engine_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(
|
|
engine_repository: Repository,
|
|
) -> None:
|
|
first = (await engine_repository.accept(command("delivery-1"))).job
|
|
second = (await engine_repository.accept(command("delivery-2"))).job
|
|
await engine_repository.apply("grant-2", PermissionGranted(job_id=second.id))
|
|
|
|
assert await engine_repository.claim_task(QueueName.JOBS) is None
|
|
|
|
await engine_repository.apply("deny-1", PermissionDenied(job_id=first.id))
|
|
task = await engine_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(
|
|
engine_repository: Repository,
|
|
) -> None:
|
|
await engine_repository.accept(command("delivery-1", issue=3))
|
|
second = (await engine_repository.accept(command("delivery-2", issue=4))).job
|
|
await engine_repository.apply("grant-2", PermissionGranted(job_id=second.id))
|
|
|
|
task = await engine_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(
|
|
engine_repository: Repository,
|
|
) -> None:
|
|
first = (await engine_repository.accept(command("delivery-1", issue=3))).job
|
|
second = (await engine_repository.accept(command("delivery-2", issue=4))).job
|
|
await engine_repository.apply("grant-2", PermissionGranted(job_id=second.id))
|
|
await engine_repository.apply("grant-1", PermissionGranted(job_id=first.id))
|
|
|
|
claimed = [
|
|
await engine_repository.claim_task(QueueName.JOBS),
|
|
await engine_repository.claim_task(QueueName.JOBS),
|
|
]
|
|
|
|
assert [task.job_id for task in claimed if task is not None] == [first.id, second.id]
|
|
assert await engine_repository.claim_task(QueueName.JOBS) is None
|
|
|
|
|
|
async def test_concurrent_claims_do_not_duplicate_task(engine_repository: Repository) -> None:
|
|
await engine_repository.accept(command("delivery-1"))
|
|
|
|
claims = await run_concurrently(
|
|
lambda: engine_repository.claim_task(QueueName.CONTROL),
|
|
lambda: engine_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 engine_repository.claim_task(QueueName.CONTROL) is None
|
|
|
|
|
|
async def test_duplicate_event_id_cannot_be_reused_for_another_job(
|
|
engine_repository: Repository,
|
|
) -> None:
|
|
first = (await engine_repository.accept(command("delivery-1", issue=3))).job
|
|
second = (await engine_repository.accept(command("delivery-2", issue=4))).job
|
|
await engine_repository.apply("permission", PermissionGranted(job_id=first.id))
|
|
|
|
with pytest.raises(RuntimeError):
|
|
await engine_repository.apply("permission", PermissionGranted(job_id=second.id))
|