1use chrono::{Duration, Utc};
 2use pretty_assertions::assert_eq;
 3
 4use crate::db::billing_subscription::SubscriptionKind;
 5use crate::db::{UserId, billing_subscription};
 6use crate::llm::db::LlmDatabase;
 7use crate::test_llm_db;
 8
 9test_llm_db!(
10    test_transfer_existing_subscription_usage,
11    test_transfer_existing_subscription_usage_postgres
12);
13
14async fn test_transfer_existing_subscription_usage(db: &mut LlmDatabase) {
15    let user_id = UserId(1);
16
17    let now = Utc::now();
18
19    let trial_period_start_at = now - Duration::days(14);
20    let trial_period_end_at = now;
21
22    let new_period_start_at = now;
23    let new_period_end_at = now + Duration::days(30);
24
25    let existing_subscription = billing_subscription::Model {
26        kind: Some(SubscriptionKind::ZedProTrial),
27        stripe_current_period_start: Some(trial_period_start_at.timestamp()),
28        stripe_current_period_end: Some(trial_period_end_at.timestamp()),
29        ..Default::default()
30    };
31
32    let existing_usage = db
33        .create_subscription_usage(
34            user_id,
35            trial_period_start_at,
36            trial_period_end_at,
37            SubscriptionKind::ZedProTrial,
38            25,
39            1_000,
40        )
41        .await
42        .unwrap();
43
44    let transferred_usage = db
45        .transfer_existing_subscription_usage(
46            user_id,
47            &existing_subscription,
48            Some(SubscriptionKind::ZedPro),
49            new_period_start_at,
50            new_period_end_at,
51        )
52        .await
53        .unwrap();
54
55    assert!(
56        transferred_usage.is_some(),
57        "subscription usage not transferred successfully"
58    );
59    let transferred_usage = transferred_usage.unwrap();
60
61    assert_eq!(
62        transferred_usage.model_requests,
63        existing_usage.model_requests
64    );
65    assert_eq!(
66        transferred_usage.edit_predictions,
67        existing_usage.edit_predictions
68    );
69}