Spaces:
Paused
Paused
| """A generation-only Studio cannot advertise or admit the other generation.""" | |
| import os | |
| import unittest | |
| from unittest.mock import patch | |
| from fastapi.testclient import TestClient | |
| from app import create_app | |
| from model_registry import PROFILES, model_registry | |
| from tetris_arena import HTTPDecisionAdapter, LOCAL_MODELS | |
| TOKEN = "t" * 32 | |
| QUESTIONS = {"accept": {"type": "noul", "instructions": "Is this accepted?"}} | |
| def generation_registry(version): | |
| return [ | |
| { | |
| "id": wire, | |
| "label": profile["repo_id"].rsplit("/", 1)[-1].split("-")[2], | |
| "version": version, | |
| "manifest_sha256": f"{index + 1:064x}", | |
| "revision": f"{index + 1:040x}", | |
| } | |
| for index, (wire, profile) in enumerate(PROFILES.items()) | |
| if f"/Decision-{version}-" in profile["repo_id"] | |
| ] | |
| class GenerationSplitTests(unittest.TestCase): | |
| def test_each_generation_is_an_independent_valid_registry(self): | |
| for version in ("1.0", "2.0"): | |
| with self.subTest(version=version): | |
| registry = model_registry(generation_registry(version)) | |
| self.assertEqual(len(registry), 6) | |
| self.assertTrue(all(f"/Decision-{version}-" in row["repo_id"] | |
| for row in registry.values())) | |
| def test_native_rejects_single_models_other_than_decision1_kai(self): | |
| for wire in ("decision2-kai", "decision-lux"): | |
| version = "2.0" if wire.startswith("decision2-") else "1.0" | |
| entry = next(row for row in generation_registry(version) if row["id"] == wire) | |
| with self.subTest(wire=wire), self.assertRaisesRegex(ValueError, "Native mode requires"): | |
| create_app(mode="native", registry=[entry]) | |
| def test_pull_queue_discovery_readiness_and_admission_follow_generation(self): | |
| for version, other in (("1.0", "2.0"), ("2.0", "1.0")): | |
| with self.subTest(version=version), patch.dict(os.environ, { | |
| "DECISION_WORKER_TOKEN": TOKEN, | |
| "TETRIS_LOCAL_API_URL": "http://127.0.0.1:7860/v1/systemone", | |
| "TETRIS_JEV_API_URL": "https://cloud.example.test/v1/systemone", | |
| "TETRIS_JEV_API_KEY": "test-cloud-key", | |
| "TETRIS_SYSTEMONE_API_URL": "https://mirror.example.test/v1/systemone", | |
| "TETRIS_SYSTEMONE_API_KEY": "test-mirror-key", | |
| }, clear=True): | |
| entries = generation_registry(version) | |
| api = create_app(mode="pull_queue", registry=entries) | |
| with TestClient(api) as client: | |
| expected = {PROFILES[row["id"]]["repo_id"] for row in entries} | |
| self.assertEqual(set(api.state.relays), {row["id"] for row in entries}) | |
| self.assertEqual({row["id"] for row in client.get("/v1/models").json()["models"]}, expected) | |
| self.assertEqual(client.get("/api/ready").status_code, 503) | |
| for index, row in enumerate(entries): | |
| heartbeat = client.post("/internal/worker/heartbeat", headers={ | |
| "Authorization": "Bearer " + TOKEN, | |
| }, json={ | |
| "model": row["id"], | |
| "manifest_sha256": row["manifest_sha256"], | |
| "worker_id": f"{index + 1:032x}", | |
| "phase": "ready", | |
| "capabilities": ["context_batch_v1"], | |
| }) | |
| self.assertEqual(heartbeat.status_code, 200, heartbeat.text) | |
| for data in ({"state": "A request"}, {"states": [{"id": "a", "state": "A request"}]}): | |
| job = client.post("/api/jobs", json={ | |
| "model": row["id"], "questions": QUESTIONS, **data, | |
| }) | |
| self.assertEqual(job.status_code, 202, job.text) | |
| self.assertEqual(job.json()["model"], row["id"]) | |
| client.delete(f"/api/jobs/{job.json()['id']}", params={"model": row["id"]}) | |
| ready = client.get("/api/ready") | |
| self.assertEqual(ready.status_code, 200) | |
| self.assertEqual(ready.json()["models"], dict.fromkeys(expected, True)) | |
| catalog = client.get("/api/tetris/config").json()["competitors"] | |
| allowed_competitors = {row.id for row in LOCAL_MODELS if row.request_model in expected} | |
| self.assertEqual({row["id"] for row in catalog if row["family"] == "decision"}, allowed_competitors) | |
| self.assertTrue(all(row["ready"] for row in catalog if row["family"] == "decision")) | |
| self.assertFalse(any(row["family"] == "cloud" for row in catalog)) | |
| self.assertNotIn("jev-cloud", api.state.tetris.adapter._endpoints) | |
| self.assertNotIn("jev-cloud-mirror", api.state.tetris.adapter._endpoints) | |
| for cloud in ("jev-cloud", "jev-cloud-mirror"): | |
| race = client.post("/api/tetris/races", json={ | |
| "left": cloud, "right": cloud, "mode": "steps", "max_steps": 1, | |
| }) | |
| self.assertEqual(race.status_code, 422, race.text) | |
| self.assertEqual(api.state.tetris.sessions, {}) | |
| excluded = generation_registry(other)[0] | |
| canonical = PROFILES[excluded["id"]]["repo_id"] | |
| for route, data in (("/v1/systemone", {"state": "A request"}), | |
| ("/v1/systemone/batches", {"states": [{"id": "a", "state": "A request"}]})): | |
| self.assertEqual(client.post(route, json={ | |
| "model": canonical, "questions": QUESTIONS, **data, | |
| }).status_code, 422) | |
| self.assertEqual(client.post("/api/jobs", json={ | |
| "model": excluded["id"], "state": "A request", "questions": QUESTIONS, | |
| }).status_code, 422) | |
| excluded_competitor = next(row.id for row in LOCAL_MODELS if row.request_model == canonical) | |
| race = client.post("/api/tetris/races", json={ | |
| "left": excluded_competitor, | |
| "right": next(iter(allowed_competitors)), | |
| "mode": "steps", "max_steps": 1, | |
| }) | |
| self.assertEqual(race.status_code, 422, race.text) | |
| self.assertEqual(api.state.tetris.sessions, {}) | |
| def test_tetris_allowlist_rejects_unknown_model(self): | |
| with self.assertRaisesRegex(ValueError, "Unknown local Decision model allowlist"): | |
| HTTPDecisionAdapter.from_environment({}, allowed_local_models={"unknown/model"}) | |
| if __name__ == "__main__": | |
| unittest.main() | |