844b02aade
Real test bodies (not stubs), driven against an in-process httpx.ASGITransport override of the gateway's get_ollama_client dependency pointing at tests/integration/mock_ollama.py. Unit (target 100% on auth/, ratelimit/, budget/): - argon2id roundtrip, wrong-key, garbage encoding, needs_rehash on param change - key format/uniqueness/prefix extraction - token counter (prompt_eval_count + eval_count, embeddings, missing-counts) - translate (OpenAI <-> Ollama for chat/completion/embeddings, streaming chunks, /v1/models list shape) - allowlist (hard-blocks, effective-set semantics across allow_all/inheritance/ empty-discovered) - discovery (parse, cache roundtrip with TTL, fail-closed, tolerates redis=None) - sliding window (allow/block/reset/per-key vs per-tenant/cost-weighted) Integration (testcontainers postgres + redis + in-process mock Ollama): - auth flow (no/malformed/wrong key all return identical sanitized 401) - proxy stream (NDJSON roundtrip, audit row's token counts match, hard-blocked endpoints uniformly 403) - openai_compat (SSE chunks, data: [DONE], non-stream shape, /v1/models) - model_discovery (allow_all sees all, default-deny sees allowed ∩ discovered, /v1/models filtered, unpermitted-but-installed = nonexistent = 403, empty cache denies even allow_all) - rate_limit (429 + Retry-After + headers; Redis down ⇒ 503, never 200) - budget (decrement + headers; pre-burned counter blocks next request) - revocation (INSERT into gateway.revocations → NOTIFY → cache evicted → 401 ≤ 1s) Includes a known-issue xfail flagging a bug in ratelimit/sliding_window.py: the per-hit ZSET member uses id(object()) which returns the same id on consecutive calls, causing same-millisecond hits to overwrite instead of stacking. To be fixed in a follow-up commit.
68 lines
2.3 KiB
Python
68 lines
2.3 KiB
Python
"""Integration tests for key revocation (SPEC §4.5, §12).
|
|
|
|
INSERT INTO ``gateway.revocations`` → AFTER-INSERT trigger fires
|
|
``pg_notify('key_revoked', key_id)`` → the gateway's revocation listener evicts
|
|
the Redis cache entry for that key prefix → the next request with that key
|
|
returns 401 within ~1 second. The key's DB status is also flipped to revoked so
|
|
a full DB re-lookup likewise denies (defense in depth).
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
|
|
import httpx
|
|
import pytest
|
|
from sqlalchemy import update
|
|
|
|
from neuronetz_gateway.db.models import ApiKey, KeyStatus
|
|
from neuronetz_gateway.db.repositories import RevocationRepository
|
|
from tests.integration.conftest import IntegrationApp, IntegrationKey
|
|
|
|
pytestmark = pytest.mark.asyncio
|
|
|
|
|
|
async def _chat(client: httpx.AsyncClient, full_key: str) -> httpx.Response:
|
|
return await client.post(
|
|
"/api/chat",
|
|
headers={"Authorization": f"Bearer {full_key}"},
|
|
json={
|
|
"model": "llama3.1:8b",
|
|
"messages": [{"role": "user", "content": "hi"}],
|
|
"stream": False,
|
|
},
|
|
)
|
|
|
|
|
|
async def test_revoked_key_rejected_within_one_second(
|
|
integration_app: IntegrationApp,
|
|
client: httpx.AsyncClient,
|
|
api_key: IntegrationKey,
|
|
) -> None:
|
|
# Warm the auth cache by making a successful request first.
|
|
ok = await _chat(client, api_key.full_key)
|
|
assert ok.status_code == 200
|
|
|
|
# Revoke: flip status + insert into the outbox (fires NOTIFY via trigger).
|
|
async with integration_app.sessionmaker() as session:
|
|
await session.execute(
|
|
update(ApiKey).where(ApiKey.id == api_key.key_id).values(status=KeyStatus.revoked)
|
|
)
|
|
await RevocationRepository(session).insert(api_key.key_id, "test")
|
|
await session.commit()
|
|
|
|
# Poll up to 1s for the listener to evict the cache; the next request must
|
|
# fail with 401 (or 403; "key revoked" is a sanitized auth failure).
|
|
deadline = 1.0
|
|
waited = 0.0
|
|
step = 0.05
|
|
while waited < deadline:
|
|
resp = await _chat(client, api_key.full_key)
|
|
if resp.status_code != 200:
|
|
assert resp.status_code in {401, 403}
|
|
assert resp.json()["error"]["request_id"]
|
|
return
|
|
await asyncio.sleep(step)
|
|
waited += step
|
|
pytest.fail(f"revoked key still accepted after {deadline}s")
|