174 lines
6.1 KiB
Python
174 lines
6.1 KiB
Python
"""Quota reservation integration tests (ticket 008).
|
|
|
|
Quota-first → credits → 402 ordering, idempotency, anniversary period
|
|
boundaries, concurrency over-grant protection.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import uuid
|
|
from datetime import UTC, datetime, timedelta
|
|
|
|
import pytest
|
|
from sqlalchemy import text
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from contract_check.core.billing.errors import NoCredits
|
|
from contract_check.core.billing.quota import (
|
|
release_document_slot,
|
|
reserve_document_slot,
|
|
)
|
|
from contract_check.core.db.session import create_session_factory # noqa: F401
|
|
|
|
pytestmark = pytest.mark.integration
|
|
|
|
|
|
def _unique_tg() -> int:
|
|
return abs(int(uuid.uuid4().int % 900000000000)) + 100000000000
|
|
|
|
|
|
async def _insert_user(session: AsyncSession, credits: int = 0) -> uuid.UUID:
|
|
result = await session.execute(
|
|
text(
|
|
"INSERT INTO users (id, telegram_id, credits_left) "
|
|
"VALUES (gen_random_uuid(), :tg, :c) RETURNING id"
|
|
),
|
|
{"tg": _unique_tg(), "c": credits},
|
|
)
|
|
return result.scalar_one()
|
|
|
|
|
|
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},
|
|
)
|
|
return result.scalar_one()
|
|
|
|
|
|
async def _grant_gift(
|
|
session: AsyncSession, user_id: uuid.UUID, plan_code: str = "lite", days: int = 30
|
|
) -> uuid.UUID:
|
|
"""Admin gift subscription (logic shared with ticket 010)."""
|
|
result = await session.execute(
|
|
text(
|
|
"INSERT INTO subscriptions "
|
|
"(id, user_id, plan_code, status, current_period_start, current_period_end) "
|
|
"VALUES (gen_random_uuid(), :u, :plan, 'active', now(), now() + make_interval(days := :days)) "
|
|
"RETURNING id"
|
|
),
|
|
{"u": user_id, "plan": plan_code, "days": days},
|
|
)
|
|
return result.scalar_one()
|
|
|
|
|
|
async def _quota_used(session: AsyncSession, subscription_id: uuid.UUID) -> int:
|
|
result = await session.execute(
|
|
text("SELECT count(*) FROM quota_usage WHERE subscription_id = :s"),
|
|
{"s": subscription_id},
|
|
)
|
|
return int(result.scalar_one())
|
|
|
|
|
|
async def test_quota_first_then_credits(
|
|
db_session: AsyncSession,
|
|
) -> None:
|
|
u = await _insert_user(db_session, credits=1)
|
|
sub_id = await _grant_gift(db_session, user_id=u, plan_code="lite") # quota 5
|
|
docs = [await _insert_doc(db_session, u) for _ in range(6)]
|
|
|
|
# 5 docs consume quota, 6th falls back to credits.
|
|
for i in range(5):
|
|
assert await reserve_document_slot(db_session, u, docs[i]) == "quota"
|
|
assert await _quota_used(db_session, sub_id) == 5
|
|
assert await reserve_document_slot(db_session, u, docs[5]) == "credits"
|
|
await db_session.commit()
|
|
|
|
# One more doc → 402 (credits exhausted).
|
|
extra = await _insert_doc(db_session, u)
|
|
with pytest.raises(NoCredits):
|
|
await reserve_document_slot(db_session, u, extra)
|
|
|
|
|
|
async def test_quota_idempotent_same_document(
|
|
db_session: AsyncSession,
|
|
) -> None:
|
|
u = await _insert_user(db_session, credits=0)
|
|
sub_id = await _grant_gift(db_session, user_id=u, plan_code="lite")
|
|
d = await _insert_doc(db_session, u)
|
|
|
|
assert await reserve_document_slot(db_session, u, d) == "quota"
|
|
assert await reserve_document_slot(db_session, u, d) == "quota"
|
|
assert await _quota_used(db_session, sub_id) == 1
|
|
await db_session.commit()
|
|
|
|
|
|
async def test_concurrency_quota_no_over_grant(
|
|
db_session: AsyncSession,
|
|
) -> None:
|
|
u = await _insert_user(db_session, credits=0)
|
|
await _grant_gift(db_session, user_id=u, plan_code="lite") # quota=5
|
|
docs = [await _insert_doc(db_session, u) for _ in range(20)]
|
|
await db_session.commit() # make rows visible to concurrent sessions
|
|
|
|
async def attempt(doc_id: uuid.UUID) -> tuple[str | Exception, uuid.UUID]:
|
|
factory = create_session_factory()
|
|
try:
|
|
async with factory() as s:
|
|
result = await reserve_document_slot(s, u, doc_id)
|
|
await s.commit()
|
|
return result, doc_id
|
|
except Exception as exc:
|
|
return exc, doc_id
|
|
|
|
results = await asyncio.gather(*(attempt(d) for d in docs))
|
|
successes = [r for r, _ in results if r == "quota"]
|
|
assert len(successes) == 5, f"over-granted: {len(successes)}"
|
|
|
|
|
|
async def test_quota_respects_period_boundary(
|
|
db_session: AsyncSession,
|
|
) -> None:
|
|
u = await _insert_user(db_session, credits=0)
|
|
now = datetime.now(tz=UTC)
|
|
start = now - timedelta(days=40)
|
|
end = now - timedelta(days=10)
|
|
result = await db_session.execute(
|
|
text(
|
|
"INSERT INTO subscriptions "
|
|
"(id, user_id, plan_code, status, current_period_start, current_period_end) "
|
|
"VALUES (gen_random_uuid(), :u, 'lite', 'active', :start, :end) "
|
|
"RETURNING id"
|
|
),
|
|
{"u": u, "start": start, "end": end},
|
|
)
|
|
sub_id = result.scalar_one()
|
|
d = await _insert_doc(db_session, u)
|
|
# Expired period → quota unavailable; user has no credits → 402.
|
|
with pytest.raises(NoCredits):
|
|
await reserve_document_slot(db_session, u, d)
|
|
assert await _quota_used(db_session, sub_id) == 0
|
|
|
|
|
|
async def test_release_document_slot_only_quota(
|
|
db_session: AsyncSession,
|
|
) -> None:
|
|
# 1) quota doc: release returns True.
|
|
u = await _insert_user(db_session, credits=0)
|
|
await _grant_gift(db_session, user_id=u, plan_code="lite")
|
|
d_quota = await _insert_doc(db_session, u)
|
|
assert await reserve_document_slot(db_session, u, d_quota) == "quota"
|
|
assert await release_document_slot(db_session, d_quota) is True
|
|
|
|
# 2) credits doc: release returns False (no quota row to delete).
|
|
u2 = await _insert_user(db_session, credits=1)
|
|
d_credit = await _insert_doc(db_session, u2)
|
|
assert await reserve_document_slot(db_session, u2, d_credit) == "credits"
|
|
assert await release_document_slot(db_session, d_credit) is False
|
|
|
|
await db_session.commit()
|