- Add stripe and cryptography to shared pyproject.toml - Add recharts, @stripe/stripe-js, stripe to portal package.json (submodule) - Add billing fields to Tenant model (stripe_customer_id, subscription_status, agent_quota, trial_ends_at) - Add budget_limit_usd to Agent model - Create TenantLlmKey and StripeEvent models in billing.py (AuditBase and Base respectively) - Create KeyEncryptionService (MultiFernet encrypt/decrypt/rotate) in crypto.py - Create compute_budget_status helper in usage.py (threshold logic: ok/warning/exceeded) - Add platform_encryption_key, stripe_, slack_oauth settings to config.py - Create Alembic migration 005 with all schema changes, RLS, grants, and composite index - All 12 tests passing (key encryption roundtrip, rotation, budget thresholds)
66 lines
2.0 KiB
Python
66 lines
2.0 KiB
Python
"""
|
|
Unit tests for budget alert threshold logic.
|
|
|
|
Tests thresholds:
|
|
- No budget limit (None) → status "ok", no alert
|
|
- Usage at 50% → status "ok"
|
|
- Usage at exactly 80% → status "warning"
|
|
- Usage at 95% → status "warning"
|
|
- Usage at exactly 100% → status "exceeded"
|
|
- Usage at 120% → status "exceeded"
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import pytest
|
|
|
|
from shared.api.usage import compute_budget_status
|
|
|
|
|
|
def test_budget_alert_no_limit() -> None:
|
|
"""Agent with no budget limit (None) → status 'ok', no alert."""
|
|
status = compute_budget_status(current_usd=500.0, budget_limit_usd=None)
|
|
assert status == "ok"
|
|
|
|
|
|
def test_budget_alert_under_threshold() -> None:
|
|
"""Usage at 50% of limit → status 'ok'."""
|
|
status = compute_budget_status(current_usd=50.0, budget_limit_usd=100.0)
|
|
assert status == "ok"
|
|
|
|
|
|
def test_budget_alert_just_below_warning() -> None:
|
|
"""Usage at 79% → still 'ok' (below 80% threshold)."""
|
|
status = compute_budget_status(current_usd=79.0, budget_limit_usd=100.0)
|
|
assert status == "ok"
|
|
|
|
|
|
def test_budget_alert_warning() -> None:
|
|
"""Usage at exactly 80% → status 'warning'."""
|
|
status = compute_budget_status(current_usd=80.0, budget_limit_usd=100.0)
|
|
assert status == "warning"
|
|
|
|
|
|
def test_budget_alert_warning_mid() -> None:
|
|
"""Usage at 95% → status 'warning'."""
|
|
status = compute_budget_status(current_usd=95.0, budget_limit_usd=100.0)
|
|
assert status == "warning"
|
|
|
|
|
|
def test_budget_alert_exceeded() -> None:
|
|
"""Usage at exactly 100% → status 'exceeded'."""
|
|
status = compute_budget_status(current_usd=100.0, budget_limit_usd=100.0)
|
|
assert status == "exceeded"
|
|
|
|
|
|
def test_budget_alert_over_limit() -> None:
|
|
"""Usage at 120% → status 'exceeded'."""
|
|
status = compute_budget_status(current_usd=120.0, budget_limit_usd=100.0)
|
|
assert status == "exceeded"
|
|
|
|
|
|
def test_budget_alert_zero_usage() -> None:
|
|
"""Zero usage with a limit → status 'ok'."""
|
|
status = compute_budget_status(current_usd=0.0, budget_limit_usd=50.0)
|
|
assert status == "ok"
|