import asyncio from urllib.parse import quote, urlsplit from uuid import uuid4 from fastapi.testclient import TestClient from httpx import Response from pwdlib import PasswordHash from pydantic import SecretStr from sqlalchemy import delete, select from ai_platform.bootstrap.app import create_app from ai_platform.infrastructure.config.settings import Settings from ai_platform.infrastructure.persistence.database import make_session_factory from ai_platform.infrastructure.persistence.models import ( AdminSession, AdminUser, TenantKey, TenantKeyAuditEvent, ) from ai_platform.platform.security import digest def expect_status(response: Response, expected: int) -> None: assert response.status_code == expected def assert_secret_absent(secret: str, body: str) -> None: if secret in body: raise AssertionError("tenant key appeared in list response") def test_admin_and_provisioner_tenant_key_lifecycle(): loaded = Settings.load() database_name = urlsplit(loaded.database_url).path.rstrip("/").rsplit("/", 1)[-1] if database_name != "ai_platform_test": raise RuntimeError("Tenant key API tests require /ai_platform_test") settings = Settings( database_url=loaded.database_url, environment="development", tenant_provisioning_secret=SecretStr("test-only-provisioner-secret"), ) sessions = make_session_factory(settings.database_url) tenant_id, admin_tenant = f"tenant/{uuid4().hex}", f"tenant-admin-{uuid4().hex}" tenant_route_id = quote(tenant_id, safe="") tenant_key_path = f"/v1/admin/tenant-keys/{tenant_route_id}" admin_email, password = f"tenant-admin-{uuid4().hex}@example.invalid", "test-only-password" app = create_app(settings) async def create_admin() -> None: async with sessions() as db: db.add( AdminUser( email=admin_email, password_hash=PasswordHash.recommended().hash(password), ) ) await db.commit() await sessions.kw["bind"].dispose() async def cleanup() -> None: async with sessions() as db: tenants = [tenant_id, admin_tenant] await db.execute( delete(TenantKeyAuditEvent).where(TenantKeyAuditEvent.tenant_id.in_(tenants)) ) await db.execute(delete(AdminSession).where(AdminSession.admin_email == admin_email)) await db.execute(delete(AdminUser).where(AdminUser.email == admin_email)) await db.execute(delete(TenantKey).where(TenantKey.tenant_id.in_(tenants))) await db.commit() await sessions.kw["bind"].dispose() asyncio.run(create_admin()) try: with TestClient(app) as client: expect_status(client.get("/v1/admin/tenant-keys"), 401) expect_status(client.post("/v1/admin/tenant-keys", json={"tenant_id": tenant_id}), 401) provision_path = "/v1/platform/tenant-keys" expect_status( client.post( provision_path, headers={"Authorization": "Bearer incorrect"}, json={"tenant_id": tenant_id}, ), 401, ) provision_headers = {"Authorization": "Bearer test-only-provisioner-secret"} expect_status( client.post(provision_path, headers=provision_headers, json={"tenant_id": " "}), 422, ) expect_status( client.post( provision_path, headers=provision_headers, json={"tenant_id": tenant_id, "unexpected": True}, ), 422, ) created = client.post( provision_path, headers=provision_headers, json={"tenant_id": tenant_id} ) expect_status(created, 201) assert created.headers["cache-control"] == "no-store" initial_key = created.json()["tenant_key"] expect_status( client.post( provision_path, headers=provision_headers, json={"tenant_id": tenant_id} ), 409, ) login = client.post( "/v1/admin/login", json={"email": admin_email, "password": password} ) expect_status(login, 200) csrf = {"x-csrf-token": login.json()["csrf_token"]} listed = client.get("/v1/admin/tenant-keys") expect_status(listed, 200) item = next(x for x in listed.json()["tenants"] if x["tenant_id"] == tenant_id) assert set(item) == {"tenant_id", "active", "created_at", "updated_at"} assert item["active"] and item["created_at"] and item["updated_at"] assert_secret_absent(initial_key, listed.text) assert "verifier" not in listed.text expect_status(client.post("/v1/admin/tenant-keys", json={"tenant_id": "no-csrf"}), 403) expect_status(client.post(f"{tenant_key_path}/rotate"), 403) expect_status(client.delete(tenant_key_path), 403) expect_status( client.post("/v1/admin/tenant-keys", headers=csrf, json={"tenant_id": tenant_id}), 409, ) run_url = f"/v1/lead-scoring/runs/{uuid4()}" def tenant_auth(key: str) -> dict[str, str]: return {"Authorization": f"Bearer {key}"} expect_status(client.get(run_url, headers=tenant_auth(initial_key)), 404) admin_created = client.post( "/v1/admin/tenant-keys", headers=csrf, json={"tenant_id": admin_tenant} ) expect_status(admin_created, 201) assert admin_created.headers["cache-control"] == "no-store" admin_key = admin_created.json()["tenant_key"] assert_secret_absent(admin_key, client.get("/v1/admin/tenant-keys").text) rotated = client.post(f"{tenant_key_path}/rotate", headers=csrf) expect_status(rotated, 200) assert rotated.headers["cache-control"] == "no-store" rotated_key = rotated.json()["tenant_key"] if rotated_key == initial_key: raise AssertionError("rotation reused the previous key") expect_status(client.get(run_url, headers=tenant_auth(initial_key)), 401) expect_status(client.get(run_url, headers=tenant_auth(rotated_key)), 404) assert_secret_absent(rotated_key, client.get("/v1/admin/tenant-keys").text) revoked = client.delete(tenant_key_path, headers=csrf) expect_status(revoked, 200) assert revoked.json()["active"] is False expect_status(client.get(run_url, headers=tenant_auth(rotated_key)), 401) reactivated = client.post(f"{tenant_key_path}/rotate", headers=csrf) expect_status(reactivated, 200) assert reactivated.headers["cache-control"] == "no-store" active_key = reactivated.json()["tenant_key"] if active_key == rotated_key: raise AssertionError("rotation reused a revoked key") expect_status(client.get(run_url, headers=tenant_auth(active_key)), 404) async def assert_audit_and_verifier() -> None: try: async with sessions() as db: row = await db.get(TenantKey, tenant_id) assert row is not None and row.active and row.verifier == digest(active_key) events = ( await db.scalars( select(TenantKeyAuditEvent) .where(TenantKeyAuditEvent.tenant_id == tenant_id) .order_by(TenantKeyAuditEvent.id) ) ).all() assert [(e.actor, e.action) for e in events] == [ ("provisioner", "create"), (admin_email, "rotate"), (admin_email, "revoke"), (admin_email, "rotate"), ] assert all(e.occurred_at is not None for e in events) created_by_admin = await db.scalar( select(TenantKeyAuditEvent).where( TenantKeyAuditEvent.tenant_id == admin_tenant ) ) assert created_by_admin is not None assert (created_by_admin.actor, created_by_admin.action) == ( admin_email, "create", ) assert {c.name for c in TenantKeyAuditEvent.__table__.columns} == { "id", "tenant_id", "actor", "action", "occurred_at", } finally: await sessions.kw["bind"].dispose() asyncio.run(assert_audit_and_verifier()) finally: asyncio.run(cleanup())