230 lines
9 KiB
Python
230 lines
9 KiB
Python
"""Unit tests for the aissurance compliance plugin.
|
|
|
|
Covers the pieces that don't require a live upstream or backend:
|
|
* config/settings resolution and URL derivation
|
|
* HMAC payload signing (determinism + tamper detection)
|
|
* buffer encrypt/decrypt round-trip + bound enforcement
|
|
* transport status-code → outcome classification
|
|
* collector field mapping against a mocked Router state
|
|
"""
|
|
import json
|
|
from types import SimpleNamespace
|
|
from unittest.mock import patch
|
|
|
|
import pytest
|
|
|
|
from compliance.settings import ComplianceSettings
|
|
from compliance import payload as payload_mod
|
|
from compliance import crypto
|
|
from compliance.buffer import BufferStore
|
|
from compliance.transport import Outcome, _classify
|
|
|
|
|
|
# ── settings ────────────────────────────────────────────────────────────────
|
|
|
|
def test_url_derivation():
|
|
s = ComplianceSettings(
|
|
server_url="https://www.aissurance.eu/api/v1/discovery/receive",
|
|
tenant_id="tid-123",
|
|
)
|
|
assert s.receive_url.endswith("/discovery/receive")
|
|
assert s.health_url == "https://www.aissurance.eu/api/v1/discovery/health/tid-123"
|
|
assert s.config_report_url == "https://www.aissurance.eu/api/v1/discovery/config/report"
|
|
|
|
|
|
def test_health_url_without_tenant():
|
|
s = ComplianceSettings(tenant_id=None)
|
|
assert s.health_url.endswith("/health/unknown")
|
|
|
|
|
|
def test_validation_errors():
|
|
assert ComplianceSettings(api_key=None).validation_errors() # missing key
|
|
assert ComplianceSettings(api_key="k", server_url="http://x").validation_errors() # not https
|
|
assert ComplianceSettings(api_key="k", polling_interval=5).validation_errors() # too fast
|
|
assert ComplianceSettings(
|
|
api_key="k", server_url="https://x/receive", polling_interval=300
|
|
).validation_errors() == []
|
|
|
|
|
|
def test_empty_string_env_becomes_none():
|
|
# config.yaml ${VAR} expansion yields "" for unset vars.
|
|
s = ComplianceSettings(api_key="", tenant_id="")
|
|
assert s.api_key is None
|
|
assert s.tenant_id is None
|
|
|
|
|
|
# ── payload signing ─────────────────────────────────────────────────────────
|
|
|
|
def _cfg(**kw):
|
|
base = dict(api_key="secret-key", tenant_id="t1", router_id="r1",
|
|
server_url="https://x/api/v1/discovery/receive")
|
|
base.update(kw)
|
|
return ComplianceSettings(**base)
|
|
|
|
|
|
def test_signature_deterministic_and_key_order_independent():
|
|
cfg = _cfg()
|
|
snap = {"timestamp": "2026-01-01T00:00:00Z", "models": [{"b": 1, "a": 2}],
|
|
"endpoints": [], "telemetry": {"x": 1}}
|
|
p1 = payload_mod.build_payload(cfg, snap)
|
|
p2 = payload_mod.build_payload(cfg, snap)
|
|
assert p1["hmac_signature"] == p2["hmac_signature"]
|
|
assert p1["hmac_signature"].startswith("sha256=")
|
|
|
|
|
|
def test_signature_detects_tampering():
|
|
cfg = _cfg()
|
|
snap = {"timestamp": "t", "models": [], "endpoints": [], "telemetry": {}}
|
|
p = payload_mod.build_payload(cfg, snap)
|
|
sig = p.pop("hmac_signature")
|
|
assert payload_mod.sign(p, cfg) == sig # untouched verifies
|
|
p["router_id"] = "evil"
|
|
assert payload_mod.sign(p, cfg) != sig # tampered does not
|
|
|
|
|
|
def test_hmac_secret_overrides_api_key():
|
|
snap = {"timestamp": "t", "models": [], "endpoints": [], "telemetry": {}}
|
|
a = payload_mod.build_payload(_cfg(), snap)
|
|
b = payload_mod.build_payload(_cfg(hmac_secret="other"), snap)
|
|
assert a["hmac_signature"] != b["hmac_signature"]
|
|
|
|
|
|
def test_split_into_batches_resigns():
|
|
cfg = _cfg()
|
|
snap = {"timestamp": "t", "models": [{"name": f"m{i}"} for i in range(5)],
|
|
"endpoints": [], "telemetry": {}}
|
|
p = payload_mod.build_payload(cfg, snap)
|
|
batches = payload_mod.split_into_batches(p, batch_size=2, cfg=cfg)
|
|
assert len(batches) == 3
|
|
assert [len(b["models"]) for b in batches] == [2, 2, 1]
|
|
for b in batches:
|
|
sig = b.pop("hmac_signature")
|
|
assert payload_mod.sign(b, cfg) == sig
|
|
|
|
|
|
# ── buffer crypto + store ───────────────────────────────────────────────────
|
|
|
|
def test_crypto_roundtrip():
|
|
ct = crypto.encrypt(b"hello", "key-1")
|
|
assert ct != b"hello"
|
|
assert crypto.decrypt(ct, "key-1") == b"hello"
|
|
assert crypto.decrypt(ct, "wrong-key") is None # wrong key -> None, no raise
|
|
|
|
|
|
def test_buffer_save_load_roundtrip(tmp_path):
|
|
cfg = _cfg(buffer_dir=str(tmp_path / "buf"))
|
|
store = BufferStore(cfg)
|
|
body = {"router_id": "r1", "models": [{"name": "x"}]}
|
|
path = store.save(body, timestamp_ms=1000)
|
|
assert path is not None and path.exists()
|
|
# raw file must not contain plaintext
|
|
assert b"router_id" not in path.read_bytes()
|
|
assert store.count() == 1
|
|
assert store.load(path) == body
|
|
|
|
|
|
def test_buffer_replay_order_and_bound(tmp_path):
|
|
cfg = _cfg(buffer_dir=str(tmp_path / "buf"), max_buffer_payloads=3)
|
|
store = BufferStore(cfg)
|
|
for ts in (5000, 1000, 3000):
|
|
store.save({"ts": ts}, timestamp_ms=ts)
|
|
pending = store.pending()
|
|
loaded_ts = [store.load(p)["ts"] for p in pending]
|
|
assert loaded_ts == [1000, 3000, 5000] # oldest-first
|
|
|
|
# exceeding the bound drops the oldest
|
|
for ts in (7000, 9000):
|
|
store.save({"ts": ts}, timestamp_ms=ts)
|
|
remaining = [store.load(p)["ts"] for p in store.pending()]
|
|
assert 1000 not in remaining
|
|
assert len(remaining) == 3
|
|
|
|
|
|
def test_buffer_drops_corrupt_file(tmp_path):
|
|
cfg = _cfg(buffer_dir=str(tmp_path / "buf"))
|
|
store = BufferStore(cfg)
|
|
p = store.save({"a": 1}, timestamp_ms=1)
|
|
p.write_bytes(b"not-a-valid-fernet-token")
|
|
assert store.load(p) is None
|
|
assert not p.exists() # corrupt file removed
|
|
|
|
|
|
# ── transport classification ────────────────────────────────────────────────
|
|
|
|
@pytest.mark.parametrize("status,expected", [
|
|
(202, Outcome.OK),
|
|
(200, Outcome.OK),
|
|
(401, Outcome.AUTH_FAILED),
|
|
(403, Outcome.AUTH_FAILED),
|
|
(429, Outcome.RATE_LIMITED),
|
|
(413, Outcome.TOO_LARGE),
|
|
(500, Outcome.RETRY),
|
|
(502, Outcome.RETRY),
|
|
])
|
|
def test_status_classification(status, expected):
|
|
assert _classify(status) == expected
|
|
|
|
|
|
# ── collector field mapping ─────────────────────────────────────────────────
|
|
|
|
async def test_collector_maps_endpoints_and_models():
|
|
from compliance import collector
|
|
|
|
cfg = SimpleNamespace(
|
|
endpoints=["http://ollama:11434"],
|
|
llama_server_endpoints=["https://llama:8000/v1"],
|
|
api_keys={"https://llama:8000/v1": "k"},
|
|
endpoint_config={},
|
|
max_concurrent_connections=4,
|
|
)
|
|
|
|
async def fake_health(ep, timeout=5):
|
|
return {"status": "ok"}
|
|
|
|
async def fake_available(ep, api_key=None):
|
|
if "ollama" in ep:
|
|
return {"mistral:7b"}
|
|
return {"org/gpt-oss-20b:Q4_K_M"}
|
|
|
|
async def fake_details(ep, route, detail):
|
|
return [{"name": "mistral:7b", "size": 4_100_000_000,
|
|
"details": {"quantization_level": "Q4_0"}}]
|
|
|
|
fake_db = SimpleNamespace(
|
|
get_last_used_map=lambda: _async_return({}),
|
|
get_token_totals_since=lambda ts: _async_return((10, 20)),
|
|
)
|
|
|
|
with patch.object(collector, "_endpoint_health", fake_health), \
|
|
patch.object(collector.fetch, "available_models", fake_available), \
|
|
patch.object(collector.fetch, "endpoint_details", fake_details), \
|
|
patch.object(collector, "get_db", lambda: fake_db), \
|
|
patch.object(collector, "get_llm_cache", lambda: None), \
|
|
patch.dict(collector.usage_counts, {"http://ollama:11434": {"mistral:7b": 2}}, clear=True), \
|
|
patch.object(collector, "get_max_connections", lambda ep: 4):
|
|
snap = await collector.collect_snapshot(cfg, started_at=0.0)
|
|
|
|
eps = {e["url"]: e for e in snap["endpoints"]}
|
|
assert eps["http://ollama:11434"]["type"] == "ollama"
|
|
assert eps["http://ollama:11434"]["tls_enabled"] is False
|
|
assert eps["http://ollama:11434"]["concurrent_connections"] == 2
|
|
assert eps["https://llama:8000/v1"]["type"] == "vllm"
|
|
assert eps["https://llama:8000/v1"]["tls_enabled"] is True
|
|
assert eps["https://llama:8000/v1"]["auth_method"] == "bearer"
|
|
|
|
models = {m["name"]: m for m in snap["models"]}
|
|
assert models["mistral:7b"]["quantization"] == "Q4_0"
|
|
assert models["mistral:7b"]["size_gb"] == pytest.approx(4.1, rel=1e-3)
|
|
assert models["gpt-oss-20b"]["quantization"] == "Q4_K_M"
|
|
# un-instrumented fields are null, not fabricated
|
|
assert models["mistral:7b"]["avg_latency_ms"] is None
|
|
assert models["mistral:7b"]["request_count_24h"] is None
|
|
|
|
tel = snap["telemetry"]
|
|
assert tel["active_connections"] == 2
|
|
assert tel["total_bytes_in_24h"] == 10
|
|
assert tel["total_requests_24h"] is None
|
|
|
|
|
|
async def _async_return(value):
|
|
return value
|