137 lines
4.8 KiB
Python
137 lines
4.8 KiB
Python
"""Credits DB integration tests (requires real Postgres via testcontainers).
|
|
|
|
Marked `integration`; not run by the default fast suite. Verifies the atomic
|
|
reserve, idempotent refund, and the race-safety of `reserve_credit`.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import uuid
|
|
|
|
import pytest
|
|
import pytest_asyncio
|
|
from sqlalchemy import text
|
|
from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine
|
|
from testcontainers.community.postgres import PostgresContainer
|
|
|
|
from contract_check.core.credits import refund_credit, reserve_credit
|
|
|
|
pytestmark = pytest.mark.integration
|
|
|
|
|
|
@pytest_asyncio.fixture
|
|
async def pg_session():
|
|
with PostgresContainer("postgres:16-alpine") as pg:
|
|
url = pg.get_connection_url().replace("psycopg2", "asyncpg")
|
|
engine = create_async_engine(url)
|
|
# Create minimal schema in-memory (Postgres handles UUID/JSONB fine).
|
|
async with engine.begin() as conn:
|
|
await conn.execute(text("CREATE EXTENSION IF NOT EXISTS pgcrypto"))
|
|
await conn.execute(
|
|
text(
|
|
"CREATE TABLE users ("
|
|
"id UUID PRIMARY KEY DEFAULT gen_random_uuid(),"
|
|
"credits_left INT NOT NULL DEFAULT 0 CHECK (credits_left >= 0)"
|
|
")"
|
|
)
|
|
)
|
|
await conn.execute(
|
|
text(
|
|
"CREATE TABLE documents ("
|
|
"id UUID PRIMARY KEY DEFAULT gen_random_uuid(),"
|
|
"user_id UUID NOT NULL REFERENCES users(id) ON DELETE CASCADE,"
|
|
"refunded BOOLEAN NOT NULL DEFAULT FALSE"
|
|
")"
|
|
)
|
|
)
|
|
async with engine.connect() as conn:
|
|
async_session = AsyncSession(bind=conn)
|
|
yield async_session
|
|
await engine.dispose()
|
|
|
|
|
|
async def insert_user(session: AsyncSession, credits: int = 5) -> uuid.UUID:
|
|
result = await session.execute(
|
|
text("INSERT INTO users (credits_left) VALUES (:c) 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 (user_id, refunded) VALUES (:u, FALSE) 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(pg_session: AsyncSession) -> None:
|
|
u = await insert_user(pg_session, credits=2)
|
|
ok = await reserve_credit(pg_session, u)
|
|
assert ok is True
|
|
await pg_session.commit()
|
|
assert await credit_balance(pg_session, u) == 1
|
|
|
|
|
|
async def test_reserve_credit_rejects_when_zero(pg_session: AsyncSession) -> None:
|
|
u = await insert_user(pg_session, credits=0)
|
|
ok = await reserve_credit(pg_session, u)
|
|
assert ok is False
|
|
|
|
|
|
async def test_reserve_credit_never_goes_negative(pg_session: AsyncSession) -> None:
|
|
u = await insert_user(pg_session, credits=1)
|
|
results = await asyncio.gather(
|
|
reserve_credit(pg_session, u),
|
|
reserve_credit(pg_session, u),
|
|
reserve_credit(pg_session, u),
|
|
)
|
|
# The atomic UPDATE WHERE credits_left>0 serializes concurrent attempts.
|
|
assert sum(1 for r in results if r) == 1
|
|
await pg_session.commit()
|
|
assert await credit_balance(pg_session, u) == 0
|
|
|
|
|
|
async def test_refund_credit_idempotent(pg_session: AsyncSession) -> None:
|
|
u = await insert_user(pg_session, credits=0)
|
|
d = await insert_doc(pg_session, u)
|
|
|
|
ok1 = await refund_credit(pg_session, d, "llm_quota", "all")
|
|
assert ok1 is True
|
|
await pg_session.commit()
|
|
assert await credit_balance(pg_session, u) == 1
|
|
|
|
ok2 = await refund_credit(pg_session, d, "llm_quota", "all")
|
|
assert ok2 is False # already refunded
|
|
await pg_session.commit()
|
|
assert await credit_balance(pg_session, u) == 1
|
|
|
|
|
|
async def test_refund_credit_respects_infra_only(pg_session: AsyncSession) -> None:
|
|
u = await insert_user(pg_session, credits=0)
|
|
d = await insert_doc(pg_session, u)
|
|
|
|
ok = await refund_credit(pg_session, d, "extraction_failed", "infra_only")
|
|
assert ok is False # user pays for garbage
|
|
await pg_session.commit()
|
|
assert await credit_balance(pg_session, u) == 0
|
|
|
|
# LLM failure is refunded under infra_only.
|
|
d2 = await insert_doc(pg_session, u)
|
|
ok2 = await refund_credit(pg_session, d2, "llm_quota", "infra_only")
|
|
assert ok2 is True
|
|
await pg_session.commit()
|
|
assert await credit_balance(pg_session, u) == 1
|