Spaces:
Running
Running
File size: 4,420 Bytes
f7318fd 6ce064b | 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 | from contextlib import contextmanager
from fastapi.testclient import TestClient
from token_trail_backend.api import create_app
from token_trail_backend.config import MAX_PROMPT_CHARS, Settings
class FakeEngine:
ready = True
def close(self):
pass
def trace(self, request, cancel=None):
return {"prompt": request.prompt, "model": request.model, "steps": []}
@contextmanager
def client(engine=None):
with TestClient(
create_app(Settings(rate_limit_per_minute=100), engine or FakeEngine(), load_model=False)
) as test_client:
yield test_client
def test_valid_trace():
with client() as api:
response = api.post("/v1/traces", json={"prompt": "Hello"})
assert response.status_code == 200
assert response.json()["prompt"] == "Hello"
def test_request_bounds():
with client() as api:
for payload in (
{"prompt": "x" * (MAX_PROMPT_CHARS + 1)},
{"prompt": "x", "max_new_tokens": 33},
{"prompt": "x", "top_k": 11},
{"prompt": "x", "model": "other"},
{"prompt": "x", "unexpected": 1},
{"prompt": "x", "max_new_tokens": "2"},
):
assert api.post("/v1/traces", json=payload).status_code in (413, 422)
def test_payload_limit_and_malformed_json():
with client() as api:
assert api.post("/v1/traces", content=b"x" * 4097).status_code == 413
assert (
api.post(
"/v1/traces",
content=b'{"prompt":"\\ud800"}',
headers={"Content-Type": "application/json"},
).status_code
== 422
)
assert (
api.post(
"/v1/traces", content=b"{", headers={"Content-Type": "application/json"}
).status_code
== 422
)
invalid_utf8 = api.post(
"/v1/traces", content=b"\xff", headers={"Content-Type": "application/json"}
)
assert invalid_utf8.status_code == 400
assert invalid_utf8.json()["error"]["code"] == "invalid_request"
def test_health_and_model_allowlist():
with client() as api:
assert api.get("/health").json() == {"status": "ready", "ready": True}
models = api.get("/v1/models").json()["models"]
assert len(models) == 1 and models[0]["alias"] == "qwen2.5-0.5b"
engine = FakeEngine()
engine.ready = False
with client(engine) as api:
assert api.get("/health").json()["ready"] is False
assert api.post("/v1/traces", json={"prompt": "x"}).status_code == 503
def test_model_load_failure_does_not_fall_back():
class BrokenEngine(FakeEngine):
ready = False
def load(self):
raise RuntimeError("pinned model unavailable")
try:
with TestClient(create_app(engine=BrokenEngine())):
pass
except RuntimeError as exc:
assert "pinned model unavailable" in str(exc)
else:
raise AssertionError("startup silently substituted a model")
def test_rate_limit():
with TestClient(
create_app(Settings(rate_limit_per_minute=1), FakeEngine(), load_model=False)
) as api:
assert api.post("/v1/traces", json={"prompt": "a"}).status_code == 200
assert api.post("/v1/traces", json={"prompt": "b"}).status_code == 429
def test_cors():
with client() as api:
allowed = api.options(
"/v1/traces",
headers={"Origin": "https://crunchycodes.net", "Access-Control-Request-Method": "POST"},
)
denied = api.options(
"/v1/traces",
headers={"Origin": "https://other.example", "Access-Control-Request-Method": "POST"},
)
assert allowed.headers.get("access-control-allow-origin") == "https://crunchycodes.net"
assert "access-control-allow-origin" not in denied.headers
def test_trace_origin_is_enforced():
with client() as api:
allowed = api.post(
"/v1/traces",
json={"prompt": "x"},
headers={"Origin": "https://crunchycodes.net"},
)
denied = api.post(
"/v1/traces",
json={"prompt": "x"},
headers={"Origin": "https://unrelated.example"},
)
assert allowed.status_code == 200
assert denied.status_code == 403
assert denied.json()["error"]["code"] == "origin_not_allowed"
|