scripts/verify_postgres_concurrent_billing.py

from __future__ import annotations

from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass
from datetime import date
from decimal import Decimal
import os
from pathlib import Path
import sys
import threading
import uuid
from unittest.mock import MagicMock, patch

PROJECT_ROOT = Path(__file__).resolve().parents[1]
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))

from fastapi import FastAPI, Request, Response
from fastapi.testclient import TestClient
from sqlalchemy import create_engine, delete, func, inspect, select
from sqlalchemy.engine import make_url
from sqlalchemy.orm import Session, sessionmaker
from starlette.middleware.base import BaseHTTPMiddleware

import app.main as app_main
from app.main import ChatResponse, chat
from app.models.database import CreditEvent, LLMModel, Tenant, get_db
from app.services.credit_service import CreditService
from app.services.dependencies import get_retrieval_service
from app.services.llm_service import LLMResult


DEFAULT_CONCURRENT_REQUESTS = 5
VERIFICATION_PROVIDER = "verification"
TENANT_NAME_PREFIX = "postgres billing verification"


@dataclass(frozen=True)
class VerificationResult:
    status_codes: list[int]
    final_balance: Decimal
    credit_event_count: int
    save_conversation_count: int


def require_postgresql_url(database_url: str) -> None:
    url = make_url(database_url)
    if not url.get_backend_name().startswith("postgresql"):
        raise ValueError(
            "PostgreSQL-backed billing verification requires a PostgreSQL "
            f"DATABASE_URL, got {url.get_backend_name()!r}."
        )


def build_verification_model_refs(run_id: str) -> list[str]:
    return [
        f"{VERIFICATION_PROVIDER}/{run_id}-primary",
        f"{VERIFICATION_PROVIDER}/{run_id}-fallback-a",
        f"{VERIFICATION_PROVIDER}/{run_id}-fallback-b",
    ]


def build_expected_status_codes(concurrent_requests: int) -> list[int]:
    if concurrent_requests < 2:
        raise ValueError("Concurrent billing verification requires at least 2 requests.")
    return [200, *([402] * (concurrent_requests - 1))]


def _split_model_ref(model_ref: str) -> tuple[str, str]:
    provider, separator, model_name = model_ref.partition("/")
    if not separator:
        raise ValueError(f"Expected provider/model format, got {model_ref!r}")
    return provider, model_name


def _verify_required_tables(engine) -> None:
    inspector = inspect(engine)
    missing = [
        table_name
        for table_name in ("tenants", "llm_models", "credit_events")
        if not inspector.has_table(table_name)
    ]
    if missing:
        raise RuntimeError(
            "Database is missing required tables for billing verification: "
            f"{', '.join(missing)}. Run migrations or start the app first."
        )


def _seed_verification_rows(
    db: Session,
    *,
    run_id: str,
    model_refs: list[str],
) -> tuple[uuid.UUID, Decimal]:
    tenant = Tenant(
        name=f"{TENANT_NAME_PREFIX} {run_id}",
        api_key=f"verification-{run_id}",
        is_active=True,
        customization={},
        balance_usd=Decimal("0"),
        billing_mode="credits",
    )
    db.add(tenant)

    for model_ref in model_refs:
        provider, model_name = _split_model_ref(model_ref)
        db.add(
            LLMModel(
                provider=provider,
                model_name=model_name,
                usd_per_1k_input=Decimal("1.000000"),
                usd_per_1k_output=Decimal("2.000000"),
                effective_from=date(2026, 1, 1),
            )
        )

    db.commit()
    db.refresh(tenant)

    expected_cost = CreditService().estimate_max_chat_cost(
        model_refs=model_refs,
        max_input_tokens=app_main.config.nlp["max_input_tokens"],
        max_output_tokens=app_main.config.chat_max_output_tokens,
        db=db,
    )
    tenant.balance_usd = expected_cost
    db.commit()

    return tenant.id, expected_cost


def _build_chat_app(session_factory, tenant_id: uuid.UUID) -> tuple[FastAPI, MagicMock]:
    class DummyTenantMiddleware(BaseHTTPMiddleware):
        async def dispatch(self, request: Request, call_next) -> Response:
            request.state.tenant_id = tenant_id
            return await call_next(request)

    def get_db_override():
        session = session_factory()
        try:
            yield session
        finally:
            session.close()

    mock_retrieval_service = MagicMock()
    mock_retrieval_service.retrieve.return_value = ([], [])

    test_app = FastAPI()
    test_app.add_middleware(DummyTenantMiddleware)
    test_app.post(f"/api/{app_main.config.api_version}/chat", response_model=ChatResponse)(
        chat
    )
    test_app.dependency_overrides[get_db] = get_db_override
    test_app.dependency_overrides[get_retrieval_service] = lambda: mock_retrieval_service

    return test_app, mock_retrieval_service


def _post_chat(test_app: FastAPI):
    with TestClient(test_app, raise_server_exceptions=False) as client:
        return client.post(
            f"/api/{app_main.config.api_version}/chat",
            json={"message": "PostgreSQL billing verification"},
        )


def wait_for_billing_barrier(barrier: threading.Barrier) -> None:
    try:
        barrier.wait(timeout=10)
    except threading.BrokenBarrierError as exc:
        raise RuntimeError(
            "Timed out waiting for both concurrent chat requests to reach "
            "the billing deduction barrier."
        ) from exc


def _audit_result(
    session_factory,
    *,
    tenant_id: uuid.UUID,
    responses,
    save_conversation_mock: MagicMock,
) -> VerificationResult:
    with session_factory() as db:
        tenant = db.get(Tenant, tenant_id)
        event_count = db.scalar(
            select(func.count())
            .select_from(CreditEvent)
            .where(CreditEvent.tenant_id == tenant_id)
        )
        if event_count is None:
            event_count = 0

        result = VerificationResult(
            status_codes=sorted(response.status_code for response in responses),
            final_balance=tenant.balance_usd,
            credit_event_count=event_count,
            save_conversation_count=save_conversation_mock.call_count,
        )

    return result


def assert_expected_result(result: VerificationResult) -> None:
    expected_status_codes = build_expected_status_codes(len(result.status_codes))
    if result.status_codes != expected_status_codes:
        raise AssertionError(
            f"Expected status split {expected_status_codes}, got {result.status_codes}"
        )
    if result.final_balance != Decimal("0.000000"):
        raise AssertionError(
            f"Expected final balance 0.000000, got {result.final_balance}"
        )
    if result.credit_event_count != 1:
        raise AssertionError(
            f"Expected exactly one CreditEvent, got {result.credit_event_count}"
        )
    if result.save_conversation_count != 1:
        raise AssertionError(
            "Expected exactly one saved conversation after the winning charge, "
            f"got {result.save_conversation_count}"
        )


def _cleanup(
    session_factory,
    *,
    tenant_id: uuid.UUID | None,
    run_id: str,
    model_refs: list[str],
) -> None:
    model_names = [_split_model_ref(model_ref)[1] for model_ref in model_refs]
    with session_factory() as db:
        if tenant_id is not None:
            db.execute(delete(CreditEvent).where(CreditEvent.tenant_id == tenant_id))
            db.execute(delete(Tenant).where(Tenant.id == tenant_id))
        else:
            db.execute(
                delete(Tenant).where(Tenant.name == f"{TENANT_NAME_PREFIX} {run_id}")
            )
        db.execute(
            delete(LLMModel).where(
                LLMModel.provider == VERIFICATION_PROVIDER,
                LLMModel.model_name.in_(model_names),
            )
        )
        db.commit()


def run_verification(
    database_url: str,
    *,
    concurrent_requests: int = DEFAULT_CONCURRENT_REQUESTS,
) -> VerificationResult:
    require_postgresql_url(database_url)
    build_expected_status_codes(concurrent_requests)
    run_id = uuid.uuid4().hex[:12]
    model_refs = build_verification_model_refs(run_id)

    engine = create_engine(database_url, pool_pre_ping=True)
    _verify_required_tables(engine)
    SessionLocal = sessionmaker(bind=engine, autocommit=False, autoflush=False)
    tenant_id: uuid.UUID | None = None

    original_primary = app_main.config.chat_model_primary
    original_fallbacks = list(app_main.config.chat_model_fallbacks)

    try:
        app_main.config.chat_model_primary = model_refs[0]
        app_main.config.chat_model_fallbacks = model_refs[1:]

        with SessionLocal() as setup_db:
            tenant_id, expected_cost = _seed_verification_rows(
                setup_db,
                run_id=run_id,
                model_refs=model_refs,
            )

        llm_result = LLMResult(
            content="Verified PostgreSQL billing race",
            model_used=model_refs[0],
            input_tokens=app_main.config.nlp["max_input_tokens"],
            output_tokens=app_main.config.chat_max_output_tokens,
        )
        test_app, _ = _build_chat_app(SessionLocal, tenant_id)
        barrier = threading.Barrier(concurrent_requests)
        save_conversation_mock = MagicMock(return_value=None)
        original_deduct = CreditService.deduct_for_chat

        def synchronized_deduct(
            self,
            tenant_id,
            amount,
            request_id,
            llm_model_id,
            db,
        ):
            wait_for_billing_barrier(barrier)
            return original_deduct(
                self,
                tenant_id,
                amount,
                request_id,
                llm_model_id,
                db,
            )

        with patch("app.main.nlp_engine") as mock_nlp, patch(
            "app.main.conversation_manager"
        ) as mock_cm, patch.object(
            CreditService,
            "deduct_for_chat",
            autospec=True,
            side_effect=synchronized_deduct,
        ):
            mock_nlp.detect_language.return_value = "en"
            mock_nlp.classify_intent.return_value = ("greeting", 0.95)
            mock_nlp.extract_entities.return_value = {}
            mock_cm.get_response.return_value = (
                None,
                "Verified PostgreSQL billing race",
                False,
                llm_result,
            )
            mock_cm.save_conversation = save_conversation_mock

            with ThreadPoolExecutor(max_workers=concurrent_requests) as executor:
                futures = [
                    executor.submit(_post_chat, test_app)
                    for _ in range(concurrent_requests)
                ]
                responses = [future.result(timeout=20) for future in futures]

        result = _audit_result(
            SessionLocal,
            tenant_id=tenant_id,
            responses=responses,
            save_conversation_mock=save_conversation_mock,
        )
        assert_expected_result(result)
        print(
            "PASS PostgreSQL concurrent billing verification: "
            f"concurrent_requests={concurrent_requests}, "
            f"status_codes={result.status_codes}, "
            f"final_balance={result.final_balance}, "
            f"credit_event_count={result.credit_event_count}, "
            f"save_conversation_count={result.save_conversation_count}, "
            f"expected_cost={expected_cost}"
        )
        return result
    finally:
        app_main.config.chat_model_primary = original_primary
        app_main.config.chat_model_fallbacks = original_fallbacks
        _cleanup(
            SessionLocal,
            tenant_id=tenant_id,
            run_id=run_id,
            model_refs=model_refs,
        )
        engine.dispose()


def main() -> int:
    database_url = os.getenv("DATABASE_URL", "")
    try:
        run_verification(database_url)
    except Exception as exc:  # pylint: disable=broad-exception-caught
        print(
            f"FAIL PostgreSQL concurrent billing verification: {exc}",
            file=sys.stderr,
        )
        return 1
    return 0


if __name__ == "__main__":
    raise SystemExit(main())