167 lines
5.1 KiB
Python
167 lines
5.1 KiB
Python
from pathlib import Path
|
|
|
|
from agentci.domain.models import Job, JobKind, JobStatus, Workflow, WorkflowKind
|
|
from agentci.worker import Worker
|
|
from agentci.workflows.common import JobRejected
|
|
|
|
|
|
class FakeStorage:
|
|
def __init__(self, workflow: Workflow | None) -> None:
|
|
self.workflow = workflow
|
|
|
|
async def get_workflow(self, _workflow_id: str) -> Workflow | None:
|
|
return self.workflow
|
|
|
|
|
|
class FakeOpenCode:
|
|
def __init__(self) -> None:
|
|
self.aborted: set[tuple[str, Path]] = set()
|
|
|
|
async def abort(self, session_id: str, workspace: Path) -> None:
|
|
self.aborted.add((session_id, workspace))
|
|
|
|
|
|
class JobStorage(FakeStorage):
|
|
def __init__(self) -> None:
|
|
super().__init__(None)
|
|
self.updates: list[tuple[JobStatus | None, str | None, str | None]] = []
|
|
self.comment_ids: list[int] = []
|
|
|
|
async def update_job(self, _job_id: str, **values) -> None:
|
|
self.updates.append(
|
|
(values.get("status"), values.get("stage"), values.get("error"))
|
|
)
|
|
|
|
async def set_job_comment(self, _job_id: str, _column: str, comment_id: int) -> None:
|
|
self.comment_ids.append(comment_id)
|
|
|
|
async def job_stage(self, _job_id: str) -> str:
|
|
return "working"
|
|
|
|
async def fail_job_workflow(self, _job_id: str) -> None:
|
|
return None
|
|
|
|
|
|
class FakeGitea:
|
|
def __init__(self) -> None:
|
|
self.created: list[str] = []
|
|
self.updated: list[tuple[int, str]] = []
|
|
|
|
async def create_comment(self, _owner: str, _repo: str, _number: int, body: str) -> int:
|
|
self.created.append(body)
|
|
return 42
|
|
|
|
async def update_comment(
|
|
self, _owner: str, _repo: str, comment_id: int, body: str
|
|
) -> None:
|
|
self.updated.append((comment_id, body))
|
|
|
|
|
|
class RejectingDispatcher:
|
|
async def dispatch(self, _job: Job) -> None:
|
|
raise JobRejected("not applicable")
|
|
|
|
|
|
def job(*, workflow_id: str | None, runtime_session_id: str | None = None) -> Job:
|
|
return Job(
|
|
id="job",
|
|
kind=JobKind.FIX,
|
|
target_key="org/repo:pr:1",
|
|
repo_owner="org",
|
|
repo_name="repo",
|
|
issue_number=1,
|
|
pr_number=1,
|
|
requester="alice",
|
|
message="",
|
|
comment_id=1,
|
|
workflow_id=workflow_id,
|
|
runtime_session_id=runtime_session_id,
|
|
)
|
|
|
|
|
|
def worker(tmp_path: Path, storage: FakeStorage, opencode: FakeOpenCode) -> Worker:
|
|
return Worker(
|
|
storage=storage, # type: ignore[arg-type]
|
|
gitea=None, # type: ignore[arg-type]
|
|
opencode=opencode, # type: ignore[arg-type]
|
|
dispatcher=None, # type: ignore[arg-type]
|
|
poll_seconds=1,
|
|
workspaces_dir=tmp_path,
|
|
)
|
|
|
|
|
|
async def test_recovery_aborts_all_workflow_sessions(tmp_path: Path) -> None:
|
|
workspace = tmp_path / "workflow" / "repo"
|
|
workflow = Workflow(
|
|
id="flow",
|
|
kind=WorkflowKind.IMPLEMENT,
|
|
repo_owner="org",
|
|
repo_name="repo",
|
|
issue_number=1,
|
|
workspace_path=workspace,
|
|
base_sha="base",
|
|
primary_session_id="primary",
|
|
reviewer_session_id="reviewer",
|
|
)
|
|
opencode = FakeOpenCode()
|
|
|
|
await worker(tmp_path, FakeStorage(workflow), opencode)._abort_job_sessions(
|
|
job(workflow_id=workflow.id)
|
|
)
|
|
|
|
assert opencode.aborted == {("primary", workspace), ("reviewer", workspace)}
|
|
|
|
|
|
async def test_recovery_aborts_one_shot_fix_session(tmp_path: Path) -> None:
|
|
opencode = FakeOpenCode()
|
|
|
|
await worker(tmp_path, FakeStorage(None), opencode)._abort_job_sessions(
|
|
job(workflow_id=None, runtime_session_id="fix-session")
|
|
)
|
|
|
|
assert opencode.aborted == {("fix-session", tmp_path / "fix-job" / "repo")}
|
|
|
|
|
|
async def test_job_status_updates_existing_gitea_comment(tmp_path: Path) -> None:
|
|
storage = JobStorage()
|
|
gitea = FakeGitea()
|
|
active_job = job(workflow_id=None)
|
|
active_job.accepted_comment_id = 41
|
|
value = Worker(
|
|
storage=storage, # type: ignore[arg-type]
|
|
gitea=gitea, # type: ignore[arg-type]
|
|
opencode=FakeOpenCode(), # type: ignore[arg-type]
|
|
dispatcher=RejectingDispatcher(), # type: ignore[arg-type]
|
|
poll_seconds=1,
|
|
workspaces_dir=tmp_path,
|
|
)
|
|
|
|
await value._run_job(active_job)
|
|
|
|
assert gitea.created == []
|
|
assert [comment_id for comment_id, _body in gitea.updated] == [41, 41]
|
|
assert "started" in gitea.updated[0][1]
|
|
assert "rejected" in gitea.updated[1][1]
|
|
assert storage.updates[-1][:2] == (JobStatus.REJECTED, "rejected")
|
|
|
|
|
|
async def test_worker_creates_only_one_comment_when_queue_comment_is_missing(
|
|
tmp_path: Path,
|
|
) -> None:
|
|
storage = JobStorage()
|
|
gitea = FakeGitea()
|
|
value = Worker(
|
|
storage=storage, # type: ignore[arg-type]
|
|
gitea=gitea, # type: ignore[arg-type]
|
|
opencode=FakeOpenCode(), # type: ignore[arg-type]
|
|
dispatcher=RejectingDispatcher(), # type: ignore[arg-type]
|
|
poll_seconds=1,
|
|
workspaces_dir=tmp_path,
|
|
)
|
|
|
|
await value._run_job(job(workflow_id=None))
|
|
|
|
assert len(gitea.created) == 1
|
|
assert storage.comment_ids == [42]
|
|
assert [comment_id for comment_id, _body in gitea.updated] == [42]
|