DataOps-platform/tests/integration/test_wp09_agent_runtime_postgres.py

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()