94 lines
2.9 KiB
Python
94 lines
2.9 KiB
Python
import asyncio
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
from typing import cast
|
|
|
|
from agentci.domain.models import Workflow, WorkflowKind
|
|
from agentci.domain.state_machine import JobState
|
|
from agentci.worker import Worker, _safe_error
|
|
|
|
|
|
class FakeStorage:
|
|
def __init__(self, workflow=None) -> None:
|
|
self.workflow = workflow
|
|
|
|
async def get_workflow(self, _workflow_id):
|
|
return self.workflow
|
|
|
|
|
|
class FakeOpenCode:
|
|
def __init__(self) -> None:
|
|
self.aborted = set()
|
|
|
|
async def abort(self, session_id, workspace):
|
|
self.aborted.add((session_id, workspace))
|
|
|
|
|
|
def worker(tmp_path: Path, storage: FakeStorage, opencode: FakeOpenCode) -> Worker:
|
|
return Worker(
|
|
storage=storage, # type: ignore[arg-type]
|
|
state_machine=SimpleNamespace(), # type: ignore[arg-type]
|
|
gitea=SimpleNamespace(), # type: ignore[arg-type]
|
|
opencode=opencode, # type: ignore[arg-type]
|
|
dispatcher=SimpleNamespace(), # type: ignore[arg-type]
|
|
poll_seconds=1,
|
|
max_concurrent_jobs=2,
|
|
workspaces_dir=tmp_path,
|
|
bot_username="agentci",
|
|
)
|
|
|
|
|
|
async def test_abort_collects_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()
|
|
state = SimpleNamespace(workflow_id="flow", runtime_session_id=None, id="job")
|
|
await worker(tmp_path, FakeStorage(workflow), opencode)._abort_job_sessions(
|
|
cast(JobState, state)
|
|
)
|
|
assert opencode.aborted == {("primary", workspace), ("reviewer", workspace)}
|
|
|
|
|
|
async def test_abort_uses_one_shot_fix_workspace(tmp_path: Path) -> None:
|
|
opencode = FakeOpenCode()
|
|
state = SimpleNamespace(workflow_id=None, runtime_session_id="session", id="job")
|
|
await worker(tmp_path, FakeStorage(), opencode)._abort_job_sessions(
|
|
cast(JobState, state)
|
|
)
|
|
assert opencode.aborted == {("session", tmp_path / "fix-job" / "repo")}
|
|
|
|
|
|
async def test_run_starts_configured_job_consumers(tmp_path: Path, monkeypatch) -> None:
|
|
value = worker(tmp_path, FakeStorage(), FakeOpenCode())
|
|
queues = []
|
|
|
|
async def recover() -> None:
|
|
pass
|
|
|
|
async def loop(queue, _stop) -> None:
|
|
queues.append(queue)
|
|
|
|
monkeypatch.setattr(value, "_recover", recover)
|
|
monkeypatch.setattr(value, "_loop", loop)
|
|
|
|
await value.run(asyncio.Event())
|
|
|
|
assert queues.count("control") == 1
|
|
assert queues.count("jobs") == 2
|
|
|
|
|
|
def test_safe_error_is_single_line_and_bounded() -> None:
|
|
value = _safe_error(RuntimeError("bad\n" + "x" * 2000))
|
|
assert "\n" not in value
|
|
assert len(value) == 1000
|