File size: 3,883 Bytes
31c8075 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 | from dataclasses import dataclass, field
from arena.codebase_archive import LocalCodebaseArchiveStore
from arena.event_bus import EventBus
from arena.run_models import CodebaseArchive, Run, RunEvent, RunRequest, RunTask
from arena.run_store import InMemoryRunStore
from arena.service import ArenaService
from arena.swarm_runtime import SwarmRunResult
from arena.validation_models import ValidationReport
@dataclass
class FakeSwarmRuntime:
started_runs: list[Run] = field(default_factory=list)
executed_runs: list[str] = field(default_factory=list)
publish_called: bool = False
def start_run(self, run: Run):
self.started_runs.append(run)
return object()
def execute_run(self, run_id: str, **kwargs) -> SwarmRunResult:
self.executed_runs.append(run_id)
return SwarmRunResult(
summary="Created the app.",
tasks=[RunTask(id="task-1", title="Build codebase", status="completed")],
events=[RunEvent(id="event-1", message="Created the app.")],
codebase_archive=CodebaseArchive(
pointer="local-snapshot://run-1",
file_count=1,
size_bytes=10,
),
)
def close_run(self, run_id: str) -> None:
pass
class FakeValidatorGraph:
def validate(self, **kwargs) -> ValidationReport:
return ValidationReport(run_id=kwargs["run_id"], passed=True, summary="ok")
def test_event_bus_publish_and_close() -> None:
bus = EventBus()
q = bus.subscribe("run-1")
bus.publish("run-1", {"kind": "event", "message": "started"})
bus.publish("run-1", {"kind": "event", "message": "working"})
bus.close("run-1")
events = []
while True:
event = q.get(timeout=1)
if event is None:
break
events.append(event)
assert len(events) == 2
assert events[0]["message"] == "started"
assert events[1]["message"] == "working"
def test_event_bus_publish_to_nonexistent_queue_is_noop() -> None:
bus = EventBus()
bus.publish("no-such-run", {"kind": "event"})
def test_event_bus_close_removes_queue() -> None:
bus = EventBus()
q = bus.subscribe("run-1")
bus.close("run-1")
assert "run-1" not in bus._queues
assert q.get(timeout=1) is None
def test_start_run_is_nonblocking_and_publishes_events() -> None:
runtime = FakeSwarmRuntime()
service = ArenaService(
run_store=InMemoryRunStore(),
swarm_runtime=runtime,
validator_graph=FakeValidatorGraph(),
archive_store=LocalCodebaseArchiveStore(),
)
access = service.start_run(RunRequest(prompt="Build a CLI", user_tests=["pytest"]))
assert access.token
assert access.run.prompt == "Build a CLI"
assert runtime.started_runs[0].id == access.run.id
def test_start_run_async_completes_and_validates() -> None:
import time
service = ArenaService(
run_store=InMemoryRunStore(),
swarm_runtime=FakeSwarmRuntime(),
validator_graph=FakeValidatorGraph(),
archive_store=LocalCodebaseArchiveStore(),
)
access = service.start_run(RunRequest(prompt="Build a CLI", user_tests=["pytest"]))
time.sleep(0.1)
view = service.get_run(access.run.id)
assert view.status in ("working", "validating", "completed")
assert len(view.tasks) == 1
assert view.tasks[0].title == "Build codebase"
assert view.codebase_archive is not None
if view.status == "completed":
assert view.validation_report is not None
def test_service_subscribe_events_returns_a_queue() -> None:
service = ArenaService(
run_store=InMemoryRunStore(),
swarm_runtime=FakeSwarmRuntime(),
validator_graph=FakeValidatorGraph(),
archive_store=LocalCodebaseArchiveStore(),
)
queue = service.subscribe_events("run-1")
assert queue is not None
|