ab24fc4c91
Pre-flight quota check for folder_index. Returns 402 with reason
when file cap or monthly token budget would be exceeded; 200 {"ok": true}
otherwise. Also adds auth_headers_free fixture to conftest.
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
95 lines
3.0 KiB
Python
95 lines
3.0 KiB
Python
"""Folder quota helpers."""
|
|
from __future__ import annotations
|
|
|
|
from datetime import datetime, timezone
|
|
|
|
import pytest
|
|
from sqlalchemy import select
|
|
|
|
from app.billing.quota import (
|
|
check_folder_quota,
|
|
add_token_usage,
|
|
QuotaExceeded,
|
|
)
|
|
from app.models import MonthlyTokenUsage
|
|
|
|
|
|
pytestmark = pytest.mark.asyncio
|
|
|
|
|
|
async def test_check_folder_quota_free_rejects_above_file_cap(db, test_user_free):
|
|
with pytest.raises(QuotaExceeded) as exc:
|
|
await check_folder_quota(
|
|
user_id=test_user_free.id, tier="free", estimated_files=500, db=db
|
|
)
|
|
assert exc.value.reason == "max_files"
|
|
|
|
|
|
async def test_check_folder_quota_free_passes_under_cap(db, test_user_free):
|
|
# No raise
|
|
await check_folder_quota(
|
|
user_id=test_user_free.id, tier="free", estimated_files=50, db=db
|
|
)
|
|
|
|
|
|
async def test_check_folder_quota_rejects_when_monthly_exhausted(db, test_user_free):
|
|
ym = datetime.now(timezone.utc).strftime("%Y-%m")
|
|
db.add(MonthlyTokenUsage(
|
|
user_id=test_user_free.id, year_month=ym, feature="folder_index", tokens_used=100_000
|
|
))
|
|
await db.commit()
|
|
with pytest.raises(QuotaExceeded) as exc:
|
|
await check_folder_quota(
|
|
user_id=test_user_free.id, tier="free", estimated_files=10, db=db
|
|
)
|
|
assert exc.value.reason == "monthly_tokens"
|
|
|
|
|
|
async def test_check_folder_quota_power_unlimited(db, test_user_power):
|
|
await check_folder_quota(
|
|
user_id=test_user_power.id, tier="power", estimated_files=999_999, db=db
|
|
)
|
|
|
|
|
|
async def test_add_token_usage_atomic_increment(db, test_user_free):
|
|
await add_token_usage(user_id=test_user_free.id, feature="folder_index", tokens=1500, db=db)
|
|
await add_token_usage(user_id=test_user_free.id, feature="folder_index", tokens=2500, db=db)
|
|
ym = datetime.now(timezone.utc).strftime("%Y-%m")
|
|
row = (await db.execute(
|
|
select(MonthlyTokenUsage).where(
|
|
MonthlyTokenUsage.user_id == test_user_free.id,
|
|
MonthlyTokenUsage.year_month == ym,
|
|
MonthlyTokenUsage.feature == "folder_index",
|
|
)
|
|
)).scalar_one()
|
|
assert row.tokens_used == 4000
|
|
|
|
|
|
async def test_add_token_usage_returns_exhausted_when_over_cap(db, test_user_free):
|
|
result = await add_token_usage(
|
|
user_id=test_user_free.id, feature="folder_index", tokens=150_000, db=db, cap=100_000
|
|
)
|
|
assert result.exhausted is True
|
|
assert result.tokens_used == 150_000
|
|
|
|
|
|
def test_quota_check_endpoint_rejects(client, auth_headers_free):
|
|
res = client.post(
|
|
"/api/v1/billing/quota/check",
|
|
json={"feature": "folder_index", "estimated_files": 500},
|
|
headers=auth_headers_free,
|
|
)
|
|
assert res.status_code == 402
|
|
body = res.json()
|
|
assert body["detail"]["reason"] == "max_files"
|
|
|
|
|
|
def test_quota_check_endpoint_passes(client, auth_headers_free):
|
|
res = client.post(
|
|
"/api/v1/billing/quota/check",
|
|
json={"feature": "folder_index", "estimated_files": 50},
|
|
headers=auth_headers_free,
|
|
)
|
|
assert res.status_code == 200
|
|
assert res.json() == {"ok": True}
|