First commit

tests/test_service_auth.py

@@ -0,0 +1,134 @@
+from enum import StrEnum
+
+import pytest
+from fastapi import APIRouter, Depends, FastAPI
+from fastapi.testclient import TestClient
+
+from service_auth import KeySource, ServiceAuth, UserId
+
+GATEWAY_KEY = "g" * 40
+SCHEDULER_KEY = "s" * 40
+PROM_KEY = "p" * 40
+
+
+class Caller(StrEnum):
+	GATEWAY = "gateway"
+	SCHEDULER = "scheduler"
+	PROMETHEUS = "prometheus"
+
+
+SOURCES = {
+	Caller.GATEWAY: KeySource("GATEWAY_KEY", "/nonexistent/gateway_key"),
+	Caller.SCHEDULER: KeySource("SCHEDULER_KEY", "/nonexistent/scheduler_key"),
+	Caller.PROMETHEUS: KeySource("PROMETHEUS_KEY", "/nonexistent/prometheus_key"),
+}
+
+
+@pytest.fixture(autouse=True)
+def keys(monkeypatch):
+	monkeypatch.setenv("GATEWAY_KEY", GATEWAY_KEY)
+	monkeypatch.setenv("SCHEDULER_KEY", SCHEDULER_KEY)
+	monkeypatch.setenv("PROMETHEUS_KEY", PROM_KEY)
+
+
+def make_app(auth: ServiceAuth, load: bool = True) -> FastAPI:
+	app = FastAPI(dependencies=[Depends(auth.authenticate)])
+	user_router = APIRouter(dependencies=[Depends(auth.require(Caller.GATEWAY))])
+	sched_router = APIRouter(dependencies=[Depends(auth.require(Caller.SCHEDULER))])
+
+	@user_router.get("/me")
+	def me(user_id: UserId):
+		return {"user_id": user_id}
+
+	@sched_router.post("/update")
+	def update():
+		return {"ok": True}
+
+	app.include_router(user_router)
+	app.include_router(sched_router)
+	if load:
+		auth.load_keys(app, Caller.GATEWAY, Caller.SCHEDULER)
+	return app
+
+
+@pytest.fixture
+def auth():
+	return ServiceAuth(SOURCES)
+
+
+@pytest.fixture
+def client(auth):
+	return TestClient(make_app(auth))
+
+
+def bearer(key):
+	return {"Authorization": f"Bearer {key}"}
+
+
+def test_no_key_is_401(client):
+	assert client.post("/update").status_code == 401
+
+
+def test_wrong_key_is_401(client):
+	assert client.post("/update", headers=bearer("x" * 40)).status_code == 401
+
+
+def test_right_caller_succeeds(client):
+	assert client.post("/update", headers=bearer(SCHEDULER_KEY)).status_code == 200
+
+
+def test_valid_key_wrong_caller_is_403(client):
+	r = client.get("/me", headers={**bearer(SCHEDULER_KEY), "X-User-Id": "1"})
+	assert r.status_code == 403
+
+
+def test_user_id_is_read_from_header(client):
+	r = client.get("/me", headers={**bearer(GATEWAY_KEY), "X-User-Id": "7"})
+	assert r.status_code == 200
+	assert r.json() == {"user_id": 7}
+
+
+@pytest.mark.parametrize("value", ["0", "-1", "abc"])
+def test_bad_user_id_is_422(client, value):
+	r = client.get("/me", headers={**bearer(GATEWAY_KEY), "X-User-Id": value})
+	assert r.status_code == 422
+
+
+def test_missing_user_id_is_422(client):
+	assert client.get("/me", headers=bearer(GATEWAY_KEY)).status_code == 422
+
+
+def test_key_not_loaded_on_this_app_is_rejected(auth):
+	# The metrics-style app only knows the Prometheus key
+	app = FastAPI(dependencies=[Depends(auth.require(Caller.PROMETHEUS))])
+
+	@app.get("/metrics")
+	def metrics():
+		return "ok"
+
+	auth.load_keys(app, Caller.PROMETHEUS)
+	c = TestClient(app)
+	assert c.get("/metrics", headers=bearer(PROM_KEY)).status_code == 200
+	assert c.get("/metrics", headers=bearer(GATEWAY_KEY)).status_code == 401
+
+
+def test_unloaded_keys_fail_closed(auth):
+	app = make_app(auth, load=False)
+	c = TestClient(app, raise_server_exceptions=False)
+	assert c.post("/update", headers=bearer(SCHEDULER_KEY)).status_code == 500
+
+
+def test_short_key_rejected_at_load(auth, monkeypatch):
+	monkeypatch.setenv("GATEWAY_KEY", "short")
+	with pytest.raises(ValueError, match="too short"):
+		auth.load_keys(FastAPI(), Caller.GATEWAY)
+
+
+def test_shared_key_rejected_at_load(auth, monkeypatch):
+	monkeypatch.setenv("SCHEDULER_KEY", GATEWAY_KEY)
+	with pytest.raises(ValueError, match="share"):
+		auth.load_keys(FastAPI(), Caller.GATEWAY, Caller.SCHEDULER)
+
+
+def test_key_paths(auth):
+	assert auth.key_paths(Caller.GATEWAY) == {"/nonexistent/gateway_key"}