clawmates/crates/cm-billing/tests/billing.rs

use cm_billing::{charge, credits_for_tokens, redeem_promo, usage_last_7_days, BillingError};
use cm_domain::{
    AccessPolicy, Agent, AgentId, AgentStatus, Role, User, UserId, Workspace, WorkspaceId,
};

#[test]
fn credit_math_rounds_up_with_a_floor_of_one() {
    assert_eq!(credits_for_tokens(1), 1);
    assert_eq!(credits_for_tokens(999), 1);
    assert_eq!(credits_for_tokens(1000), 1);
    assert_eq!(credits_for_tokens(1001), 2);
    assert_eq!(credits_for_tokens(0), 1, "a charged run costs at least 1");
}

async fn seeded(pool: &sqlx::PgPool) -> (Workspace, Agent, uuid::Uuid) {
    let ws = Workspace {
        id: WorkspaceId::new(),
        name: "Acme".into(),
        plan: "team".into(),
    };
    cm_db::repo::workspaces::insert(pool, &ws).await.unwrap();
    let owner = User {
        id: UserId::new(),
        workspace_id: ws.id,
        email: format!("{}@acme.test", UserId::new()),
        role: Role::Owner,
        display_name: "Owner".into(),
        created_at: time::OffsetDateTime::UNIX_EPOCH,
    };
    cm_db::repo::users::insert(pool, &owner).await.unwrap();
    let agent = Agent {
        id: AgentId::new(),
        workspace_id: ws.id,
        name: "Scout".into(),
        job_title: "Analyst".into(),
        system_prompt: String::new(),
        avatar: String::new(),
        accent: String::new(),
        wallpaper: String::new(),
        managed_by: owner.id,
        status: AgentStatus::Online,
    };
    cm_db::repo::agents::insert(pool, &agent, &AccessPolicy::default())
        .await
        .unwrap();
    let session = cm_db::repo::sessions::create(pool, agent.id, ws.id, "Chat")
        .await
        .unwrap();
    let run_id = cm_db::repo::runs::create(pool, session.id).await.unwrap();
    (ws, agent, run_id)
}

#[tokio::test]
async fn charges_span_lots_oldest_first_and_record_usage() {
    let pool = cm_testkit::test_pool().await;
    let (ws, agent, run_id) = seeded(&pool).await;
    cm_db::repo::credits::add_lot(&pool, ws.id, 2, "first")
        .await
        .unwrap();
    cm_db::repo::credits::add_lot(&pool, ws.id, 100, "second")
        .await
        .unwrap();

    // 2500 tokens → 3 credits: drains the first lot (2) then one more.
    let deducted = charge(&pool, ws.id, agent.id, run_id, 1500, 1000)
        .await
        .unwrap();
    assert_eq!(deducted, 3);
    assert_eq!(
        cm_db::repo::credits::balance(&pool, ws.id).await.unwrap(),
        99
    );

    let remaining: Vec<i64> = sqlx::query_scalar::<_, sqlx::types::BigDecimal>(
        "SELECT remaining FROM credit_lots WHERE workspace_id = $1 ORDER BY purchased_at, id",
    )
    .bind(ws.id.as_uuid())
    .fetch_all(&pool)
    .await
    .unwrap()
    .into_iter()
    .map(|d| d.with_scale(0).to_string().parse().unwrap())
    .collect();
    assert_eq!(remaining, vec![0, 99], "oldest lot drains first");

    let (tin, tout, credits) = usage_last_7_days(&pool, ws.id).await.unwrap();
    assert_eq!((tin, tout, credits), (1500, 1000, 3));
}

#[tokio::test]
async fn an_empty_workspace_records_usage_but_clamps_at_zero() {
    let pool = cm_testkit::test_pool().await;
    let (ws, agent, run_id) = seeded(&pool).await;
    cm_db::repo::credits::add_lot(&pool, ws.id, 1, "tiny")
        .await
        .unwrap();

    // Owes 5, only 1 available: deducts 1, balance hits zero, never negative.
    let deducted = charge(&pool, ws.id, agent.id, run_id, 4000, 500)
        .await
        .unwrap();
    assert_eq!(deducted, 1);
    assert_eq!(
        cm_db::repo::credits::balance(&pool, ws.id).await.unwrap(),
        0
    );

    // The full obligation is still on the books.
    let (_, _, credits) = usage_last_7_days(&pool, ws.id).await.unwrap();
    assert_eq!(credits, 5);
}

#[tokio::test]
async fn promo_codes_redeem_exactly_once() {
    let pool = cm_testkit::test_pool().await;
    let (ws, _, _) = seeded(&pool).await;
    sqlx::query("INSERT INTO promo_codes (code, credits) VALUES ('WELCOME500', 500)")
        .execute(&pool)
        .await
        .unwrap();

    let granted = redeem_promo(&pool, ws.id, "WELCOME500").await.unwrap();
    assert_eq!(granted, 500);
    assert_eq!(
        cm_db::repo::credits::balance(&pool, ws.id).await.unwrap(),
        500
    );

    let again = redeem_promo(&pool, ws.id, "WELCOME500").await;
    assert!(matches!(again, Err(BillingError::PromoUnavailable)));
    let unknown = redeem_promo(&pool, ws.id, "NOPE").await;
    assert!(matches!(unknown, Err(BillingError::PromoUnavailable)));
}