Add test suite and switch API to lifespan handler
test_entropy.py covers the HMAC-DRBG (determinism, reseed, length,
reseed limit), health checks (random passes, frozen/all-ones/short/
long-run rejected), a bit-balance smoke test on DRBG output, and the
API endpoints including auth gating. All 16 tests pass without a
camera (producer is stubbed).
Switch api.py from the deprecated @app.on_event('startup') to the
lifespan async context manager. Add pytest to requirements.txt.
Generated by Mistral Vibe.
Co-Authored-By: Mistral Vibe <vibe@mistral.ai>
This commit is contained in:
@@ -13,6 +13,7 @@ Endpoints:
|
||||
import os
|
||||
import threading
|
||||
import time
|
||||
from contextlib import asynccontextmanager
|
||||
from datetime import datetime, timezone
|
||||
|
||||
from fastapi import FastAPI, Header, HTTPException
|
||||
@@ -29,8 +30,6 @@ API_KEY = os.environ.get("ENTROPY_API_KEY", "")
|
||||
|
||||
NUM_BYTES = BITS // 8
|
||||
|
||||
app = FastAPI(title="Entropy-RNG Service", version="1.0")
|
||||
|
||||
|
||||
class LatestBlob:
|
||||
"""Thread-safe holder for the most recent entropy blob."""
|
||||
@@ -76,10 +75,14 @@ def producer_loop():
|
||||
time.sleep(INTERVAL)
|
||||
|
||||
|
||||
@app.on_event("startup")
|
||||
def _start_producer():
|
||||
@asynccontextmanager
|
||||
async def lifespan(app):
|
||||
t = threading.Thread(target=producer_loop, daemon=True)
|
||||
t.start()
|
||||
yield
|
||||
|
||||
|
||||
app = FastAPI(title="Entropy-RNG Service", version="1.0", lifespan=lifespan)
|
||||
|
||||
|
||||
class EntropyResponse(BaseModel):
|
||||
|
||||
@@ -3,3 +3,5 @@ opencv-python>=4.8,<5
|
||||
# FastAPI stack for the entropy service (added in the API stage)
|
||||
fastapi>=0.110
|
||||
uvicorn[standard]>=0.29
|
||||
# Test dependencies
|
||||
pytest>=8.0
|
||||
|
||||
+146
@@ -0,0 +1,146 @@
|
||||
"""Tests for the entropy source: DRBG, health checks, and API.
|
||||
|
||||
Run with: pytest -q
|
||||
These tests stub the camera so they run without hardware.
|
||||
"""
|
||||
import os
|
||||
|
||||
import pytest
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
import entropy
|
||||
import health
|
||||
|
||||
|
||||
# --- HMAC-DRBG ---
|
||||
|
||||
def test_drbg_deterministic_same_seed():
|
||||
a = entropy.HMACDRBG(b"seed-A").generate(32)
|
||||
b = entropy.HMACDRBG(b"seed-A").generate(32)
|
||||
assert a == b
|
||||
|
||||
|
||||
def test_drbg_different_seed_different_output():
|
||||
a = entropy.HMACDRBG(b"seed-A").generate(32)
|
||||
b = entropy.HMACDRBG(b"seed-B").generate(32)
|
||||
assert a != b
|
||||
|
||||
|
||||
def test_drbg_reseed_changes_output():
|
||||
d = entropy.HMACDRBG(b"seed")
|
||||
first = d.generate(32)
|
||||
d.reseed(b"new-entropy")
|
||||
second = d.generate(32)
|
||||
assert first != second
|
||||
|
||||
|
||||
def test_drbg_output_length():
|
||||
d = entropy.HMACDRBG(b"seed")
|
||||
assert len(d.generate(1)) == 1
|
||||
assert len(d.generate(64)) == 64
|
||||
assert len(d.generate(1000)) == 1000
|
||||
|
||||
|
||||
def test_drbg_requires_reseed_after_limit():
|
||||
d = entropy.HMACDRBG(b"seed")
|
||||
for _ in range(10001):
|
||||
d.generate(1)
|
||||
with pytest.raises(RuntimeError, match="reseed"):
|
||||
d.generate(1)
|
||||
|
||||
|
||||
# --- Health checks ---
|
||||
|
||||
def test_health_passes_random_data():
|
||||
health.check(os.urandom(2500))
|
||||
|
||||
|
||||
def test_health_rejects_frozen_frame():
|
||||
with pytest.raises(health.HealthCheckError):
|
||||
health.check(b"\x00" * 2500)
|
||||
|
||||
|
||||
def test_health_rejects_all_ones():
|
||||
with pytest.raises(health.HealthCheckError):
|
||||
health.check(b"\xff" * 2500)
|
||||
|
||||
|
||||
def test_health_rejects_short_sample():
|
||||
with pytest.raises(health.HealthCheckError, match="too short"):
|
||||
health.check(os.urandom(100))
|
||||
|
||||
|
||||
def test_health_long_run_rejected():
|
||||
# 2500 bytes with a single long run of zeros embedded
|
||||
data = bytearray(os.urandom(2500))
|
||||
data[1000:1000 + 30] = b"\x00" * 30 # 240-bit run of zeros
|
||||
with pytest.raises(health.HealthCheckError, match="long run"):
|
||||
health.check(bytes(data))
|
||||
|
||||
|
||||
# --- Statistical smoke test on DRBG output ---
|
||||
|
||||
def test_drbg_output_is_balanced():
|
||||
"""Generated bits should be roughly 50/50 over a large sample."""
|
||||
d = entropy.HMACDRBG(b"seed")
|
||||
out = d.generate(2500)
|
||||
ones = sum(bin(byte).count("1") for byte in out)
|
||||
total = len(out) * 8
|
||||
# Expect ~50%; allow 45-55% to avoid flakiness
|
||||
ratio = ones / total
|
||||
assert 0.45 <= ratio <= 0.55, f"bit balance {ratio:.3f} outside [0.45, 0.55]"
|
||||
|
||||
|
||||
# --- API ---
|
||||
|
||||
@pytest.fixture
|
||||
def client(monkeypatch):
|
||||
# Disable the real producer thread; we inject blobs manually
|
||||
monkeypatch.setattr("api.producer_loop", lambda: None)
|
||||
import api
|
||||
api.latest.set(b"\x00" * 32)
|
||||
return TestClient(api.app)
|
||||
|
||||
|
||||
def test_api_healthz(client):
|
||||
r = client.get("/healthz")
|
||||
assert r.status_code == 200
|
||||
assert r.json() == {"status": "ok"}
|
||||
|
||||
|
||||
def test_api_root(client):
|
||||
r = client.get("/")
|
||||
assert r.status_code == 200
|
||||
body = r.json()
|
||||
assert body["service"] == "entropy-rng"
|
||||
assert body["bits"] == 256
|
||||
|
||||
|
||||
def test_api_entropy_returns_blob(client):
|
||||
r = client.get("/entropy")
|
||||
assert r.status_code == 200
|
||||
body = r.json()
|
||||
assert body["hex"] == "00" * 32
|
||||
assert body["bits"] == 256
|
||||
assert "ts" in body
|
||||
|
||||
|
||||
def test_api_entropy_requires_key_when_set(client, monkeypatch):
|
||||
import api
|
||||
monkeypatch.setattr(api, "API_KEY", "secret")
|
||||
# No key
|
||||
r = client.get("/entropy")
|
||||
assert r.status_code == 401
|
||||
# Wrong key
|
||||
r = client.get("/entropy", headers={"X-API-Key": "wrong"})
|
||||
assert r.status_code == 401
|
||||
# Right key
|
||||
r = client.get("/entropy", headers={"X-API-Key": "secret"})
|
||||
assert r.status_code == 200
|
||||
|
||||
|
||||
def test_api_healthz_open_even_with_auth(client, monkeypatch):
|
||||
import api
|
||||
monkeypatch.setattr(api, "API_KEY", "secret")
|
||||
r = client.get("/healthz")
|
||||
assert r.status_code == 200
|
||||
Reference in New Issue
Block a user