from __future__ import annotations import os import uuid import pytest from sqlalchemy import create_engine, text from app.core.llm.runtime_governance import BudgetPolicy from app.core.llm.runtime_governance_repository import ( SqlAlchemyRuntimeGovernanceRepository, ) pytestmark = pytest.mark.integration def test_wp09_two_connections_allow_only_one_budget_lease_and_append_only_audit(): database_url = os.getenv("TEST_DATABASE_URL") if not database_url: pytest.skip("TEST_DATABASE_URL is required") tenant_id = f"wp09_{uuid.uuid4().hex[:16]}" engine = create_engine(database_url) try: with engine.begin() as connection: repository = SqlAlchemyRuntimeGovernanceRepository(connection) repository.create_budget(tenant_id, BudgetPolicy(5, 10, 1, 100, 1, {"controlled-model-v1": 1})) first = engine.connect() second = engine.connect() transaction_one = first.begin() transaction_two = second.begin() try: repo_one = SqlAlchemyRuntimeGovernanceRepository(first) repo_two = SqlAlchemyRuntimeGovernanceRepository(second) assert repo_one.reserve_budget(tenant_id, "lease-a", "controlled-model-v1", 5, 10, 1, 50) > 0 transaction_one.commit() with pytest.raises(ValueError, match="budget_exhausted"): repo_two.reserve_budget(tenant_id, "lease-b", "controlled-model-v1", 1, 1, 1, 1) transaction_two.rollback() with engine.begin() as connection: repository = SqlAlchemyRuntimeGovernanceRepository(connection) repository.append_audit({"invocation_id": str(uuid.uuid4()), "tenant_id": tenant_id, "idempotency_key": "audit-a", "route_id": "route-a", "generation": "generation-1", "lease_fence": 1, "input_hash": "a" * 64, "evidence_digests": ["b" * 64], "decision": "authorized"}) with pytest.raises(ValueError, match="append_only_conflict"): repository.append_audit({"invocation_id": str(uuid.uuid4()), "tenant_id": tenant_id, "idempotency_key": "audit-a", "route_id": "route-a", "generation": "generation-1", "lease_fence": 1, "input_hash": "a" * 64, "evidence_digests": ["b" * 64], "decision": "authorized"}) finally: if transaction_one.is_active: transaction_one.rollback() if transaction_two.is_active: transaction_two.rollback() first.close() second.close() finally: with engine.begin() as connection: connection.execute(text("DELETE FROM public.agent_invocation_audits WHERE tenant_id = :tenant_id"), {"tenant_id": tenant_id}) connection.execute(text("DELETE FROM public.agent_runtime_reservations WHERE tenant_id = :tenant_id"), {"tenant_id": tenant_id}) connection.execute(text("DELETE FROM public.agent_runtime_budgets WHERE tenant_id = :tenant_id"), {"tenant_id": tenant_id}) engine.dispose()