"""Credits DB integration tests (requires running Docker Compose infrastructure). Marked `integration`; not run by the default fast suite. Verifies the atomic reserve, idempotent refund, and the race-safety of `reserve_credit` against the real Postgres schema. """ from __future__ import annotations import asyncio import uuid import pytest from sqlalchemy import text from sqlalchemy.ext.asyncio import AsyncSession from contract_check.core.credits import refund_credit, reserve_credit pytestmark = pytest.mark.integration async def insert_user(session: AsyncSession, credits: int = 5) -> uuid.UUID: result = await session.execute( text( "INSERT INTO users (id, email, password_hash, credits_left, is_active) " "VALUES (gen_random_uuid(), gen_random_uuid() || '@test.local', '', :c, TRUE) " "RETURNING id" ), {"c": credits}, ) user_id = result.scalar_one() await session.commit() return user_id async def insert_doc(session: AsyncSession, user_id: uuid.UUID) -> uuid.UUID: result = await session.execute( text( "INSERT INTO documents (id, user_id, s3_key, filename, mime, bytes, status) " "VALUES (gen_random_uuid(), :u, 's3://test', 'test.pdf', 'application/pdf', 0, 'queued') " "RETURNING id" ), {"u": user_id}, ) doc_id = result.scalar_one() await session.commit() return doc_id async def credit_balance(session: AsyncSession, user_id: uuid.UUID) -> int: result = await session.execute( text("SELECT credits_left FROM users WHERE id = :u"), {"u": user_id} ) return int(result.scalar_one()) async def test_reserve_credit_decrements_once(db_session: AsyncSession) -> None: u = await insert_user(db_session, credits=2) ok = await reserve_credit(db_session, u) assert ok is True await db_session.commit() assert await credit_balance(db_session, u) == 1 async def test_reserve_credit_rejects_when_zero(db_session: AsyncSession) -> None: u = await insert_user(db_session, credits=0) ok = await reserve_credit(db_session, u) assert ok is False async def test_reserve_credit_never_goes_negative(db_session: AsyncSession) -> None: u = await insert_user(db_session, credits=1) async def _attempt() -> bool: # Each concurrent attempt must use its own DB session/connection. async with AsyncSession(db_session.bind) as session: result = await reserve_credit(session, u) await session.commit() return result results = await asyncio.gather(_attempt(), _attempt(), _attempt()) # The atomic UPDATE WHERE credits_left>0 serializes concurrent attempts. assert sum(1 for r in results if r) == 1 assert await credit_balance(db_session, u) == 0 async def test_refund_credit_idempotent(db_session: AsyncSession) -> None: u = await insert_user(db_session, credits=0) d = await insert_doc(db_session, u) ok1 = await refund_credit(db_session, d, "llm_quota", "all") assert ok1 is True await db_session.commit() assert await credit_balance(db_session, u) == 1 ok2 = await refund_credit(db_session, d, "llm_quota", "all") assert ok2 is False # already refunded await db_session.commit() assert await credit_balance(db_session, u) == 1 async def test_refund_credit_respects_infra_only(db_session: AsyncSession) -> None: u = await insert_user(db_session, credits=0) d = await insert_doc(db_session, u) ok = await refund_credit(db_session, d, "extraction_failed", "infra_only") assert ok is False # user pays for garbage await db_session.commit() assert await credit_balance(db_session, u) == 0 # LLM failure is refunded under infra_only. d2 = await insert_doc(db_session, u) ok2 = await refund_credit(db_session, d2, "llm_quota", "infra_only") assert ok2 is True await db_session.commit() assert await credit_balance(db_session, u) == 1 async def event_count(db_session: AsyncSession, user_id: uuid.UUID, kind: str) -> int: result = await db_session.execute( text("SELECT count(*) FROM credit_events WHERE user_id = :u AND kind = :k"), {"u": user_id, "k": kind}, ) return int(result.scalar_one()) async def test_reserve_emits_reserve_event(db_session: AsyncSession) -> None: u = await insert_user(db_session, credits=3) ok = await reserve_credit(db_session, u) assert ok is True await db_session.commit() assert await event_count(db_session, u, "reserve") == 1 async def test_refund_emits_refund_auto_event(db_session: AsyncSession) -> None: u = await insert_user(db_session, credits=1) d = await insert_doc(db_session, u) await reserve_credit(db_session, u, document_id=d) await db_session.commit() await refund_credit(db_session, d, "llm_quota", "all") await db_session.commit() assert await event_count(db_session, u, "reserve") == 1 assert await event_count(db_session, u, "refund_auto") == 1 async def test_adjust_emits_admin_grant_event(db_session: AsyncSession) -> None: from contract_check.core.credits import adjust_credits u = await insert_user(db_session, credits=0) bal = await adjust_credits(db_session, u, 10) await db_session.commit() assert bal == 10 assert await event_count(db_session, u, "admin_grant") == 1 # idempotent zero delta emits nothing bal2 = await adjust_credits(db_session, u, 0) await db_session.commit() assert bal2 == 10 assert await event_count(db_session, u, "admin_grant") == 1 async def test_balance_after_chain_matches_final_balance(db_session: AsyncSession) -> None: u = await insert_user(db_session, credits=5) d1 = await insert_doc(db_session, u) d2 = await insert_doc(db_session, u) d3 = await insert_doc(db_session, u) await reserve_credit(db_session, u, document_id=d1) await reserve_credit(db_session, u, document_id=d2) await reserve_credit(db_session, u, document_id=d3) await refund_credit(db_session, d2, "llm_quota", "all") from contract_check.core.credits import adjust_credits await adjust_credits(db_session, u, 10) await db_session.commit() result = await db_session.execute( text("SELECT balance_after FROM credit_events WHERE user_id = :u ORDER BY id DESC LIMIT 1"), {"u": u}, ) last_balance_after = int(result.scalar_one()) assert last_balance_after == await credit_balance(db_session, u)