nxtgauge-backend-rust/crates/db/tests/ai_credits.rs

//! Integration tests for the Ask Ash AI credits wallet (Phase 1).
//! Requires a live Postgres reachable via TEST_DATABASE_URL with the
//! `20260317000000_init_config_schema`, `20260317000001_init_users_schema`,
//! and `20260703210000_ai_credits_wallet` migrations already applied.
//! Run manually: `TEST_DATABASE_URL=... cargo test -p db --test ai_credits`.

use db::models::ai_credits::{AiCreditsError, AiCreditsRepository};
use sqlx::PgPool;
use uuid::Uuid;

async fn pool() -> PgPool {
    let url = std::env::var("TEST_DATABASE_URL")
        .expect("set TEST_DATABASE_URL to run ai_credits integration tests");
    sqlx::postgres::PgPoolOptions::new()
        .max_connections(10)
        .connect(&url)
        .await
        .expect("connect to test db")
}

async fn make_user(pool: &PgPool) -> Uuid {
    let id = Uuid::new_v4();
    sqlx::query("INSERT INTO users (id, email, password_hash) VALUES ($1, $2, 'x')")
        .bind(id)
        .bind(format!("{id}@test.local"))
        .execute(pool)
        .await
        .expect("insert test user");
    id
}

#[tokio::test]
async fn ensure_wallet_grants_free_plan_credits() {
    let pool = pool().await;
    let user_id = make_user(&pool).await;

    let wallet = AiCreditsRepository::ensure_wallet(&pool, user_id)
        .await
        .expect("ensure_wallet");

    assert_eq!(wallet.monthly_credits_total, 10);
    assert_eq!(wallet.available_credits(), 10);

    // Idempotent: calling again returns the same wallet, doesn't grant twice.
    let wallet2 = AiCreditsRepository::ensure_wallet(&pool, user_id)
        .await
        .expect("ensure_wallet again");
    assert_eq!(wallet2.id, wallet.id);
    assert_eq!(wallet2.monthly_credits_total, 10);
}

#[tokio::test]
async fn reserve_capture_debits_wallet() {
    let pool = pool().await;
    let user_id = make_user(&pool).await;
    AiCreditsRepository::ensure_wallet(&pool, user_id).await.unwrap();

    let hold = AiCreditsRepository::try_reserve_credits(&pool, user_id, "test_feature", 3, None, None)
        .await
        .expect("reserve");

    let mid_wallet = AiCreditsRepository::get_wallet(&pool, user_id).await.unwrap().unwrap();
    assert_eq!(mid_wallet.reserved_credits, 3);
    assert_eq!(mid_wallet.available_credits(), 7);

    let captured = AiCreditsRepository::try_capture_reservation(&pool, hold.id, None)
        .await
        .expect("capture");
    assert!(captured);

    let final_wallet = AiCreditsRepository::get_wallet(&pool, user_id).await.unwrap().unwrap();
    assert_eq!(final_wallet.reserved_credits, 0);
    assert_eq!(final_wallet.monthly_credits_used, 3);
    assert_eq!(final_wallet.lifetime_used_credits, 3);
    assert_eq!(final_wallet.available_credits(), 7);
}

#[tokio::test]
async fn reserve_release_refunds_wallet() {
    let pool = pool().await;
    let user_id = make_user(&pool).await;
    AiCreditsRepository::ensure_wallet(&pool, user_id).await.unwrap();

    let hold = AiCreditsRepository::try_reserve_credits(&pool, user_id, "test_feature", 4, None, None)
        .await
        .expect("reserve");

    let released = AiCreditsRepository::try_release_reservation(&pool, hold.id)
        .await
        .expect("release");
    assert!(released);

    let wallet = AiCreditsRepository::get_wallet(&pool, user_id).await.unwrap().unwrap();
    assert_eq!(wallet.reserved_credits, 0);
    assert_eq!(wallet.monthly_credits_used, 0);
    assert_eq!(wallet.available_credits(), 10);
    // Regression: releasing a hold must also roll back the daily action/
    // credit counters bumped at reserve time -- a failed request (e.g. the
    // LLM call errors) must not permanently burn daily quota when the user
    // was never actually charged.
    assert_eq!(wallet.daily_actions_used, 0, "release must roll back daily_actions_used");
    assert_eq!(wallet.daily_credits_used, 0, "release must roll back daily_credits_used");

    // Releasing an already-released hold is a no-op, not an error.
    let released_again = AiCreditsRepository::try_release_reservation(&pool, hold.id)
        .await
        .expect("release again");
    assert!(!released_again);
}

#[tokio::test]
async fn reserve_beyond_balance_is_rejected() {
    let pool = pool().await;
    let user_id = make_user(&pool).await;
    AiCreditsRepository::ensure_wallet(&pool, user_id).await.unwrap();

    let result = AiCreditsRepository::try_reserve_credits(&pool, user_id, "test_feature", 11, None, None).await;
    // The seeded Free plan's daily_credit_limit (10) equals its
    // monthly_credits_total (10), so a request for 11 trips the
    // daily-limit guardrail before the balance check ever runs -- both
    // are correct rejections of the same over-request, so accept either.
    assert!(matches!(
        result,
        Err(AiCreditsError::InsufficientCredits) | Err(AiCreditsError::DailyCreditLimitReached)
    ));

    // Balance must be untouched after a rejected reservation.
    let wallet = AiCreditsRepository::get_wallet(&pool, user_id).await.unwrap().unwrap();
    assert_eq!(wallet.reserved_credits, 0);
    assert_eq!(wallet.available_credits(), 10);
}

#[tokio::test]
async fn reserve_is_idempotent_on_key() {
    let pool = pool().await;
    let user_id = make_user(&pool).await;
    AiCreditsRepository::ensure_wallet(&pool, user_id).await.unwrap();

    let key = format!("idem-{}", Uuid::new_v4());
    let hold1 = AiCreditsRepository::try_reserve_credits(&pool, user_id, "test_feature", 5, None, Some(&key))
        .await
        .expect("first reserve");
    let hold2 = AiCreditsRepository::try_reserve_credits(&pool, user_id, "test_feature", 5, None, Some(&key))
        .await
        .expect("retried reserve with same key");

    assert_eq!(hold1.id, hold2.id, "retry with same idempotency key must not create a second hold");

    let wallet = AiCreditsRepository::get_wallet(&pool, user_id).await.unwrap().unwrap();
    assert_eq!(wallet.reserved_credits, 5, "credits must only be reserved once despite the retry");
}

/// The concrete race condition this design exists to prevent (Section 4.2
/// / G4 of the architecture doc): fire N concurrent reservations against a
/// wallet that can only afford one of them, and confirm exactly one wins.
#[tokio::test]
async fn concurrent_reservations_cannot_overdraw_wallet() {
    let pool = pool().await;
    let user_id = make_user(&pool).await;
    let wallet = AiCreditsRepository::ensure_wallet(&pool, user_id).await.unwrap();
    assert_eq!(wallet.available_credits(), 10);

    // 5 concurrent requests each asking for all 10 credits. At most one
    // should succeed; a broken check-then-write implementation would let
    // several through and drive the balance negative.
    let mut handles = Vec::new();
    for _ in 0..5 {
        let pool = pool.clone();
        handles.push(tokio::spawn(async move {
            AiCreditsRepository::try_reserve_credits(&pool, user_id, "test_feature", 10, None, None).await
        }));
    }

    let mut successes = 0;
    for h in handles {
        if h.await.unwrap().is_ok() {
            successes += 1;
        }
    }

    assert_eq!(successes, 1, "exactly one concurrent reservation for the full balance should succeed");

    let final_wallet = AiCreditsRepository::get_wallet(&pool, user_id).await.unwrap().unwrap();
    assert_eq!(final_wallet.available_credits(), 0);
    assert!(
        final_wallet.monthly_credits_used <= final_wallet.monthly_credits_total,
        "wallet must never be driven into overdraft"
    );
}

#[tokio::test]
async fn daily_action_limit_enforced() {
    let pool = pool().await;
    let user_id = make_user(&pool).await;
    AiCreditsRepository::ensure_wallet(&pool, user_id).await.unwrap();

    // Free plan's daily_action_limit is 3 (seeded in the migration).
    for _ in 0..3 {
        AiCreditsRepository::try_reserve_credits(&pool, user_id, "test_feature", 1, None, None)
            .await
            .expect("reserve within daily action limit");
    }

    let fourth = AiCreditsRepository::try_reserve_credits(&pool, user_id, "test_feature", 1, None, None).await;
    assert!(matches!(fourth, Err(AiCreditsError::DailyActionLimitReached)));
}