//! 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))); }