Spaces:
Running
Running
File size: 4,167 Bytes
53d0a86 | 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 | """The Kai wire-ID change does not create a mixed-generation queue alias."""
import unittest
from fastapi.testclient import TestClient
from app import create_app
from contract import MODEL
from model_registry import PROFILES
from relay import Relay
OLD_WIRE_ID = "decision-nano-preview"
CANONICAL = PROFILES[MODEL]["repo_id"]
MANIFEST = "a" * 64
TOKEN = "t" * 32
WORKER_ID = "1" * 32
class NoTetris:
def public_config(self):
return {"competitors": []}
def payload(model):
return {
"model": model,
"state": "A duplicate payment needs a refund.",
"questions": {"refund": {
"type": "noul",
"instructions": "Is a refund requested?",
}},
}
class KaiWireIDMigrationTests(unittest.TestCase):
def setUp(self):
self.app_context = TestClient(create_app(
mode="pull_queue",
relay=Relay(TOKEN, MANIFEST, model=MODEL),
tetris_manager=NoTetris(),
registry=[{
"id": MODEL,
"label": "Kai",
"version": "1.0",
"manifest_sha256": MANIFEST,
}],
))
self.client = self.app_context.__enter__()
self.worker_headers = {"Authorization": "Bearer " + TOKEN}
def tearDown(self):
self.app_context.__exit__(None, None, None)
def test_studio_uses_new_wire_id_but_public_contract_stays_canonical(self):
models = self.client.get("/v1/models").json()
self.assertEqual(models["models"][0]["id"], CANONICAL)
self.assertEqual(models["models"][0]["wire_id"], MODEL)
self.assertEqual(models["default"], CANONICAL)
self.assertEqual(self.client.get("/api/status", params={"model": MODEL}).status_code, 200)
self.assertEqual(
self.client.get("/api/status", params={"model": OLD_WIRE_ID}).status_code,
422,
)
for route in ("/api/evaluate", "/api/jobs", "/v1/systemone"):
with self.subTest(route=route):
self.assertEqual(self.client.post(route, json=payload(OLD_WIRE_ID)).status_code, 422)
# Canonical public ID is accepted by parsing, even while no resident
# worker is connected; the unavailable queue returns 503, not 422.
self.assertEqual(
self.client.post("/v1/systemone", json=payload(CANONICAL)).status_code,
503,
)
def test_previous_worker_cannot_join_claim_or_complete_new_queue(self):
old_identity = {
"model": OLD_WIRE_ID,
"manifest_sha256": MANIFEST,
"worker_id": WORKER_ID,
}
new_identity = dict(old_identity, model=MODEL)
self.assertEqual(self.client.post(
"/internal/worker/heartbeat",
json=dict(old_identity, phase="ready"),
headers=self.worker_headers,
).status_code, 422)
self.assertEqual(self.client.post(
"/internal/worker/heartbeat",
json=dict(new_identity, phase="ready"),
headers=self.worker_headers,
).status_code, 200)
job = self.client.post("/api/jobs", json=payload(MODEL))
self.assertEqual(job.status_code, 202)
self.assertEqual(job.json()["model"], MODEL)
self.assertEqual(self.client.post(
"/internal/worker/claim",
json=old_identity,
headers=self.worker_headers,
).status_code, 422)
claim = self.client.post(
"/internal/worker/claim",
json=dict(new_identity, wait_seconds=0),
headers=self.worker_headers,
)
self.assertEqual(claim.status_code, 200)
leased = claim.json()["job"]
self.assertEqual(leased["id"], job.json()["id"])
self.assertEqual(self.client.post(
"/internal/worker/result",
json=dict(
old_identity,
id=leased["id"],
lease_token=leased["lease_token"],
error_code="input_rejected",
),
headers=self.worker_headers,
).status_code, 422)
if __name__ == "__main__":
unittest.main()
|