116 lines
3.9 KiB
Python
116 lines
3.9 KiB
Python
"""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
|