81 lines
2.3 KiB
Python
81 lines
2.3 KiB
Python
from pathlib import Path
|
|
|
|
from agentci.domain.models import Job, JobKind, Workflow, WorkflowKind
|
|
from agentci.worker import Worker
|
|
|
|
|
|
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))
|
|
|
|
|
|
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")}
|