tests/unit/test_token_budget.py

"""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