135 lines
3.6 KiB
Python
135 lines
3.6 KiB
Python
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"}
|