tests_python/test_run_flow.py

import pytest

from arena.codebase_archive import LocalCodebaseArchiveStore
from arena.event_bus import EventBus
from arena.models import Criteria
from arena.run_flow import (
    RunArchiveNotFound,
    RunFlow,
)
from arena.run_models import CodebaseArchive, RunEvent, RunRequest, RunTask
from arena.run_store import InMemoryRunStore
from arena.swarm_runtime import SwarmRuntime, SwarmRunResult
from arena.validator_graph import DeterministicValidatorExecutor, RubricReview, ValidatorGraph
from arena.validator_plan import validation_checks_from_user_tests

from conftest import make_local_session


def make_flow(archive_store=None, validator_graph=None):
    store = archive_store or LocalCodebaseArchiveStore()
    event_bus = EventBus()
    return RunFlow(
        run_store=InMemoryRunStore(),
        swarm_runtime=SwarmRuntime(
            session_factory=make_local_session,
            archive_store=store,
            event_bus=event_bus,
        ),
        validator_graph=validator_graph or ValidatorGraph(
            executor=DeterministicValidatorExecutor(),
            rubric_enabled=False,
        ),
        archive_store=store,
        event_bus=event_bus,
    )


def test_run_flow_creates_run_and_saves_artifacts() -> None:
    archive_store = LocalCodebaseArchiveStore()
    flow = make_flow(archive_store)

    access = flow.create_run(RunRequest(prompt="Build a CLI", user_tests=["echo ok"]))

    view = flow.get_run(access.run.id)
    assert access.run.status == "completed"
    assert view.status == "completed"
    assert view.tasks[0].status == "completed"
    assert "Built the app" in view.summary
    assert view.codebase_archive is not None
    assert view.codebase_archive.pointer.startswith("local-snapshot://")
    assert view.validation_checks == validation_checks_from_user_tests(["echo ok"])
    assert view.validation_report is not None


def test_run_flow_validates_run_and_marks_completed() -> None:
    archive_store = LocalCodebaseArchiveStore()
    flow = make_flow(archive_store)

    access = flow.create_run(RunRequest(prompt="Build a CLI", user_tests=["echo ok"]))
    report = flow.validate_run(access.run.id)
    view = flow.get_run(access.run.id)

    assert report.passed is True
    assert report.run_id == access.run.id
    assert view.status == "completed"
    assert view.validation_report is not None


def test_run_flow_zip_download_contains_workspace_files() -> None:
    flow = make_flow()
    access = flow.create_run(RunRequest(prompt="Build a CLI", user_tests=[]))

    payload = flow.get_codebase_zip(access.run.id)
    assert payload.filename == f"{access.run.id}.zip"
    assert payload.content_type == "application/zip"
    assert payload.data.startswith(b"PK")


def test_run_flow_requires_archive_before_download_or_validation() -> None:
    flow = make_flow()
    access = flow.run_store.create_run(RunRequest(prompt="Build"))

    with pytest.raises(RunArchiveNotFound):
        flow.get_codebase_zip(access.run.id)
    with pytest.raises(RunArchiveNotFound):
        flow.validate_run(access.run.id)


def test_run_flow_rejoin_restores_access() -> None:
    flow = make_flow()
    access = flow.create_run(RunRequest(prompt="Build a CLI", user_tests=[]))

    rejoined = flow.rejoin(access.token)

    assert rejoined.run.id == access.run.id
    assert rejoined.token == access.token


def test_run_flow_cancel_run_is_noop_for_completed_runs() -> None:
    flow = make_flow()
    access = flow.create_run(RunRequest(prompt="Build a CLI", user_tests=[]))

    cancelled = flow.cancel_run(access.run.id)

    assert cancelled.status == "completed"
    assert flow.get_run(access.run.id).status == "completed"


def test_run_flow_cancel_run_closes_swarm_runtime() -> None:
    flow = make_flow()
    access = flow.run_store.create_run(RunRequest(prompt="Build a CLI", user_tests=[]))
    closed: list[str] = []
    flow.swarm_runtime.close_run = lambda run_id: closed.append(run_id)

    cancelled = flow.cancel_run(access.run.id)

    assert closed == [access.run.id]
    assert cancelled.status == "interrupted"


def test_run_flow_add_run_message_touches_runtime_and_records_event() -> None:
    flow = make_flow()
    access = flow.create_run(RunRequest(prompt="Build a CLI", user_tests=[]))
    touched: list[str] = []
    flow.swarm_runtime.touch_run = lambda run_id: touched.append(run_id) or True

    view = flow.add_run_message(access.run.id, "More context")

    assert touched == [access.run.id]
    assert view.events[-1].message == "User: More context"


def test_run_flow_start_run_streams_late_lifecycle_events() -> None:
    import time

    flow = make_flow()
    access = flow.start_run(RunRequest(prompt="Build a CLI", user_tests=["echo ok"]))

    queue = flow.subscribe_events(access.run.id)

    deadline = time.time() + 5.0
    while time.time() < deadline and flow.get_run(access.run.id).status != "completed":
        time.sleep(0.05)

    events = _drain(queue, timeout=1.0)
    kinds = [event["kind"] for event in events]

    assert "validating" in kinds
    assert "validated" in kinds


def test_run_flow_persists_runtime_events_while_swarm_is_running() -> None:
    store = InMemoryRunStore()
    archive_store = LocalCodebaseArchiveStore()
    event_bus = EventBus()
    runtime = PublishingRuntime(event_bus=event_bus, run_store=store)
    flow = RunFlow(
        run_store=store,
        swarm_runtime=runtime,
        validator_graph=FakeValidatorGraph(),
        archive_store=archive_store,
        event_bus=event_bus,
    )

    access = flow.create_run(RunRequest(prompt="Build a CLI", user_tests=[]))

    assert runtime.live_view is not None
    assert runtime.live_view.tasks[0].status == "active"
    assert runtime.live_view.events[0].message == "Coordinator started work"
    view = flow.get_run(access.run.id)
    assert [event.id for event in view.events] == ["event-live"]


def test_run_flow_subscribe_events_before_start_run_misses_early_events() -> None:
    """Behavior contract: events fire-and-forget; subscribing after a run finishes gets nothing."""
    flow = make_flow()
    access = flow.create_run(RunRequest(prompt="Build a CLI", user_tests=["echo ok"]))

    queue = flow.subscribe_events(access.run.id)
    events = _drain(queue, timeout=0.2)

    assert events == []


def test_run_flow_validate_run_returns_cached_report_on_second_call() -> None:
    flow = make_flow()
    access = flow.create_run(RunRequest(prompt="Build a CLI", user_tests=["echo ok"]))

    first = flow.validate_run(access.run.id)
    second = flow.validate_run(access.run.id)

    assert first is second


def test_run_flow_passes_prompt_and_criteria_to_rubric_reviewer() -> None:
    captured: dict = {}

    class CapturingReviewer:
        def review(self, *, prompt, criteria, check_results):
            captured["prompt"] = prompt
            captured["criteria"] = list(criteria)
            return RubricReview(score=1.0, feedback="ok", risks=[])

    flow = make_flow(
        validator_graph=ValidatorGraph(
            executor=DeterministicValidatorExecutor(),
            rubric_reviewer=CapturingReviewer(),
            rubric_enabled=True,
        )
    )

    flow.create_run(
        RunRequest(
            prompt="Build a CLI",
            criteria=[Criteria(text="pytest must pass")],
            user_tests=["echo ok"],
        )
    )

    assert captured["prompt"] == "Build a CLI"
    assert len(captured["criteria"]) == 1
    assert captured["criteria"][0].text == "pytest must pass"


def test_run_flow_skips_rubric_when_validator_graph_has_it_disabled() -> None:
    captured: dict = {"calls": 0}

    class CountingReviewer:
        def review(self, *, prompt, criteria, check_results):
            captured["calls"] += 1
            return RubricReview(score=1.0, feedback="ok", risks=[])

    flow = make_flow(
        validator_graph=ValidatorGraph(
            executor=DeterministicValidatorExecutor(),
            rubric_reviewer=CountingReviewer(),
            rubric_enabled=False,
        )
    )

    report = flow.create_run(
        RunRequest(prompt="Build a CLI", user_tests=["echo ok"])
    ).run.validation_report

    assert captured["calls"] == 0
    assert report is not None
    assert report.rubric_score is None


def _drain(queue, *, timeout: float) -> list[dict]:
    import queue as queue_module

    events: list[dict] = []
    while True:
        try:
            event = queue.get(timeout=timeout)
        except queue_module.Empty:
            break
        if event is None:
            break
        events.append(event)
    return events


class FakeValidatorGraph:
    def validate(self, **kwargs):
        from arena.validation_models import ValidationReport

        return ValidationReport(run_id=kwargs["run_id"], passed=True, summary="ok")


class PublishingRuntime:
    def __init__(self, *, event_bus: EventBus, run_store: InMemoryRunStore) -> None:
        self.event_bus = event_bus
        self.run_store = run_store
        self.live_view = None

    def start_run(self, run):
        return object()

    def execute_run(self, run_id: str, **kwargs) -> SwarmRunResult:
        task = RunTask(id="task-live", title="Build codebase", status="active")
        event = RunEvent(id="event-live", message="Coordinator started work")
        self.event_bus.publish(run_id, {"kind": "task_created", "task": task.model_dump()})
        self.event_bus.publish(run_id, {"kind": "event", "event": event.model_dump(mode="json")})
        self.live_view = self.run_store.get_run_view(run_id)
        return SwarmRunResult(
            summary="Created the app.",
            tasks=[task.model_copy(update={"status": "completed"})],
            events=[event],
            codebase_archive=CodebaseArchive(
                pointer="local-snapshot://run-1",
                file_count=1,
                size_bytes=10,
            ),
        )

    def close_run(self, run_id: str) -> None:
        pass

    def touch_run(self, run_id: str) -> bool:
        return True