tests_python/test_run_store.py

from datetime import timedelta

import pytest

from arena.models import (
    CodebaseArchive,
    Criteria,
    ModelProviderConfig,
    RunEvent,
    RunRequest,
    RunTask,
    ValidationCheck,
    ValidationReport,
    utc_now,
)
from arena.run_store import InMemoryRunStore, InvalidRunToken, RedisRunStore, create_run_store


class FakeRedisClient:
    def __init__(self) -> None:
        self.data: dict[str, str] = {}
        self.commands: list[tuple[str, ...]] = []

    def execute(self, *args: str):
        self.commands.append(tuple(args))
        command = args[0].upper()
        if command == "GET":
            return self.data.get(args[1])
        if command == "SET":
            self.data[args[1]] = args[2]
            return "OK"
        raise AssertionError(f"unexpected command: {args}")


def test_create_run_returns_rejoin_token_and_public_view_without_secret() -> None:
    store = InMemoryRunStore()

    access = store.create_run(
        RunRequest(
            prompt="Build a small FastAPI app",
            criteria=[Criteria(text="pytest must pass")],
            user_tests=["uv run pytest"],
        )
    )

    assert access.token
    assert access.run.status == "queued"
    assert access.run.prompt == "Build a small FastAPI app"
    assert access.run.criteria[0].text == "pytest must pass"
    assert access.run.user_tests == ["uv run pytest"]
    assert access.token not in access.run.model_dump_json()


def test_create_run_persists_model_provider_without_api_key() -> None:
    store = InMemoryRunStore()

    access = store.create_run(
        RunRequest(
            prompt="Build a CLI",
            model_provider=ModelProviderConfig(
                provider="openrouter",
                model="qwen/model",
                api_key="sk-user-secret",
                base_url="https://openrouter.ai/api/v1",
            ),
        )
    )

    payload = access.run.model_dump_json()
    assert access.run.model_provider is not None
    assert access.run.model_provider.model == "qwen/model"
    assert "sk-user-secret" not in payload


def test_rejoin_run_restores_access_by_token() -> None:
    store = InMemoryRunStore()
    access = store.create_run(RunRequest(prompt="Build a CLI"))

    rejoined = store.rejoin(access.token)

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


def test_rejoin_run_rejects_unknown_token() -> None:
    store = InMemoryRunStore()

    with pytest.raises(InvalidRunToken):
        store.rejoin("not-a-real-token")


def test_store_appends_events_and_saves_task_state() -> None:
    store = InMemoryRunStore()
    access = store.create_run(RunRequest(prompt="Build a CLI"))
    created_at = access.run.updated_at
    earlier = utc_now()
    later = earlier + timedelta(seconds=1)

    store.save_tasks(
        access.run.id,
        [RunTask(id="task-1", title="Implement CLI", status="active")],
    )
    store.append_event(access.run.id, RunEvent(id="event-2", message="Coding", created_at=later))
    store.append_event(access.run.id, RunEvent(id="event-1", message="Planning", created_at=earlier))

    view = store.get_run_view(access.run.id)

    assert view.tasks[0].status == "active"
    assert [event.id for event in view.events] == ["event-1", "event-2"]
    assert view.updated_at > created_at


def test_store_deduplicates_appended_events_by_id() -> None:
    store = InMemoryRunStore()
    access = store.create_run(RunRequest(prompt="Build a CLI"))
    event = RunEvent(id="event-1", message="Planning")

    store.append_event(access.run.id, event)
    store.append_event(access.run.id, event)

    view = store.get_run_view(access.run.id)

    assert [item.id for item in view.events] == ["event-1"]


def test_store_status_mutation_updates_timestamp() -> None:
    store = InMemoryRunStore()
    access = store.create_run(RunRequest(prompt="Build a CLI"))

    view = store.save_status(access.run.id, "working")

    assert view.status == "working"
    assert view.updated_at > access.run.updated_at


def test_store_saves_archive_checks_and_validation_report() -> None:
    store = InMemoryRunStore()
    access = store.create_run(RunRequest(prompt="Build a CLI"))

    store.save_codebase_archive(
        access.run.id,
        CodebaseArchive(pointer="local-archive://run-1", file_count=2, size_bytes=256),
    )
    store.save_validation_checks(
        access.run.id,
        [
            ValidationCheck(
                id="check-1",
                name="tests",
                command="uv run pytest",
                source="generated",
            )
        ],
    )
    view = store.save_validation_report(
        access.run.id,
        ValidationReport(run_id=access.run.id, passed=True, summary="ok"),
    )

    assert view.codebase_archive is not None
    assert view.codebase_archive.pointer == "local-archive://run-1"
    assert view.validation_checks[0].name == "tests"
    assert view.validation_report is not None
    assert view.validation_report.passed is True


def test_redis_run_store_restores_run_after_restart() -> None:
    client = FakeRedisClient()
    store = RedisRunStore(client=client)
    access = store.create_run(RunRequest(prompt="Build a CLI"))

    store.save_tasks(
        access.run.id,
        [RunTask(id="task-1", title="Implement CLI", status="completed")],
    )
    store.append_event(access.run.id, RunEvent(id="event-1", message="Done"))

    restarted = RedisRunStore(client=client)
    rejoined = restarted.rejoin(access.token)

    assert rejoined.run.id == access.run.id
    assert rejoined.run.tasks[0].status == "completed"
    assert rejoined.run.events[0].message == "Done"
    assert any(command[0] == "SET" and command[1].startswith("arena:run:") for command in client.commands)


def test_store_marks_active_tasks_interrupted() -> None:
    store = InMemoryRunStore()
    access = store.create_run(RunRequest(prompt="Build a CLI"))
    store.save_tasks(
        access.run.id,
        [
            RunTask(id="task-1", title="Implement CLI", status="active"),
            RunTask(id="task-2", title="Write docs", status="completed"),
        ],
    )

    view = store.mark_active_tasks_interrupted(access.run.id)

    assert [task.status for task in view.tasks] == ["interrupted", "completed"]


def test_create_run_store_defaults_to_memory(monkeypatch: pytest.MonkeyPatch) -> None:
    monkeypatch.delenv("RUN_STORE_PROVIDER", raising=False)

    assert isinstance(create_run_store(), InMemoryRunStore)


def test_create_run_store_requires_redis_env(monkeypatch: pytest.MonkeyPatch) -> None:
    monkeypatch.setenv("RUN_STORE_PROVIDER", "redis")
    monkeypatch.delenv("UPSTASH_REDIS_REST_URL", raising=False)
    monkeypatch.delenv("UPSTASH_REDIS_REST_TOKEN", raising=False)

    with pytest.raises(ValueError):
        create_run_store()