tests_python/test_event_streaming.py

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