"""Tests: Token budget — math, atomic update, warning threshold.""" import pytest from unittest.mock import AsyncMock, MagicMock, patch def test_total_token_budget_includes_bonus(): from app.models.client_access import ClientAccess access = ClientAccess.__new__(ClientAccess) access.token_budget = 500_000 access.bonus_tokens = 50_000 assert access.total_token_budget == 550_000 def test_total_token_budget_no_bonus(): from app.models.client_access import ClientAccess access = ClientAccess.__new__(ClientAccess) access.token_budget = 200_000 access.bonus_tokens = 0 assert access.total_token_budget == 200_000 @pytest.mark.asyncio async def test_budget_exhausted_blocks_chat(client, test_author, db_session): """A token-exhausted author should get a friendly message, not an error.""" from app.models.client_access import ClientAccess from sqlalchemy import select result = await db_session.execute( select(ClientAccess).where(ClientAccess.author_id == test_author.id) ) access = result.scalar_one_or_none() if access: access.tokens_used = access.token_budget # Exhaust budget await db_session.commit() # The chat endpoint should handle exhaustion gracefully res = await client.post(f"/api/chat/{test_author.id}", json={"message": "hello", "session_id": "test-session"}) # Should not be 500 — either 200 with friendly message or 403 assert res.status_code in (200, 403, 422) def test_warning_threshold_calculation(): """80% threshold triggers warning, 95% triggers alert.""" from app.config import get_settings cfg = get_settings() budget = 500_000 used_80 = int(budget * cfg.TOKEN_WARNING_THRESHOLD_PCT) used_95 = int(budget * cfg.TOKEN_ALERT_THRESHOLD_PCT) assert used_80 == 400_000 assert used_95 == 475_000