Vapora/crates/vapora-llm-router/tests/budget_test.rs

use std::collections::HashMap;
use vapora_llm_router::{BudgetManager, RoleBudget};

fn create_test_budgets() -> HashMap<String, RoleBudget> {
    let mut budgets = HashMap::new();

    budgets.insert(
        "architect".to_string(),
        RoleBudget {
            role: "architect".to_string(),
            monthly_limit_cents: 50000, // $500
            weekly_limit_cents: 12500,  // $125
            fallback_provider: "gemini".to_string(),
            alert_threshold: 0.8,
        },
    );

    budgets.insert(
        "developer".to_string(),
        RoleBudget {
            role: "developer".to_string(),
            monthly_limit_cents: 30000,
            weekly_limit_cents: 7500,
            fallback_provider: "ollama".to_string(),
            alert_threshold: 0.8,
        },
    );

    budgets
}

#[tokio::test]
async fn test_budget_initialization() {
    let budgets = create_test_budgets();
    let manager = BudgetManager::new(budgets);

    let status = manager.check_budget("architect").await.unwrap();
    assert_eq!(status.role, "architect");
    assert_eq!(status.monthly_remaining_cents, 50000);
    assert_eq!(status.monthly_utilization, 0.0);
    assert!(!status.exceeded);
    assert!(!status.near_threshold);
}

#[tokio::test]
async fn test_budget_spending() {
    let budgets = create_test_budgets();
    let manager = BudgetManager::new(budgets);

    manager.record_spend("developer", 3000).await.unwrap();

    let status = manager.check_budget("developer").await.unwrap();
    assert_eq!(status.monthly_remaining_cents, 27000);
    assert!((status.monthly_utilization - 0.1).abs() < 0.01);
}

#[tokio::test]
async fn test_multiple_spends_accumulate() {
    let budgets = create_test_budgets();
    let manager = BudgetManager::new(budgets);

    manager.record_spend("developer", 5000).await.unwrap();
    manager.record_spend("developer", 3000).await.unwrap();
    manager.record_spend("developer", 2000).await.unwrap();

    let status = manager.check_budget("developer").await.unwrap();
    assert_eq!(status.monthly_remaining_cents, 20000); // 30000 - 10000
}

#[tokio::test]
async fn test_alert_threshold_near() {
    let budgets = create_test_budgets();
    let manager = BudgetManager::new(budgets);

    // Spend 81% of weekly budget (12500 * 0.81 = 10125) to trigger near_threshold
    // This keeps us under both monthly and weekly limits while triggering alert
    let spend_amount = (12500.0 * 0.81) as u32; // 10125
    manager
        .record_spend("architect", spend_amount)
        .await
        .unwrap();

    let status = manager.check_budget("architect").await.unwrap();
    assert!(!status.exceeded);
    assert!(status.near_threshold);
}

#[tokio::test]
async fn test_budget_exceeded() {
    let budgets = create_test_budgets();
    let manager = BudgetManager::new(budgets);

    // Spend entire monthly budget
    manager.record_spend("developer", 30000).await.unwrap();

    let status = manager.check_budget("developer").await.unwrap();
    assert!(status.exceeded);
    assert_eq!(status.monthly_remaining_cents, 0);
}

#[tokio::test]
async fn test_budget_overspend() {
    let budgets = create_test_budgets();
    let manager = BudgetManager::new(budgets);

    // Spend more than budget (overflow protection)
    manager.record_spend("developer", 35000).await.unwrap();

    let status = manager.check_budget("developer").await.unwrap();
    assert!(status.exceeded);
    assert_eq!(status.monthly_remaining_cents, 0); // Saturating subtract
}

#[tokio::test]
async fn test_weekly_budget_independent() {
    let budgets = create_test_budgets();
    let manager = BudgetManager::new(budgets);

    // Spend 100% of weekly budget but only 25% of monthly
    manager.record_spend("developer", 7500).await.unwrap();

    let status = manager.check_budget("developer").await.unwrap();
    assert_eq!(status.monthly_remaining_cents, 22500);
    assert_eq!(status.weekly_remaining_cents, 0);
    assert!(status.exceeded); // Both budgets checked
}

#[tokio::test]
async fn test_fallback_provider() {
    let budgets = create_test_budgets();
    let manager = BudgetManager::new(budgets);

    let fallback_dev = manager.get_fallback_provider("developer").await.unwrap();
    assert_eq!(fallback_dev, "ollama");

    let fallback_arch = manager.get_fallback_provider("architect").await.unwrap();
    assert_eq!(fallback_arch, "gemini");
}

#[tokio::test]
async fn test_unknown_role_error() {
    let budgets = create_test_budgets();
    let manager = BudgetManager::new(budgets);

    let result = manager.check_budget("unknown").await;
    assert!(result.is_err());

    let result = manager.record_spend("unknown", 100).await;
    assert!(result.is_err());
}

#[tokio::test]
async fn test_get_all_budgets() {
    let budgets = create_test_budgets();
    let manager = BudgetManager::new(budgets);

    manager.record_spend("architect", 5000).await.unwrap();
    manager.record_spend("developer", 3000).await.unwrap();

    let all_statuses = manager.get_all_budgets().await;
    assert_eq!(all_statuses.len(), 2);

    let arch_status = all_statuses.iter().find(|s| s.role == "architect").unwrap();
    assert_eq!(arch_status.monthly_remaining_cents, 45000);

    let dev_status = all_statuses.iter().find(|s| s.role == "developer").unwrap();
    assert_eq!(dev_status.monthly_remaining_cents, 27000);
}

#[tokio::test]
async fn test_budget_status_comprehensive() {
    let budgets = create_test_budgets();
    let manager = BudgetManager::new(budgets);

    // Spend 6000 cents: keeps us at 12% of monthly and 48% of weekly (both safe)
    manager.record_spend("architect", 6000).await.unwrap();

    let status = manager.check_budget("architect").await.unwrap();
    assert_eq!(status.monthly_remaining_cents, 44000);
    assert!((status.monthly_utilization - 0.12).abs() < 0.01);
    assert!(!status.exceeded);
    assert!(!status.near_threshold); // 12% < 80%
    assert_eq!(status.fallback_provider, "gemini");
}