File size: 6,851 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 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 | 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()
|