1use chrono::{Duration, Utc};
  2use pretty_assertions::assert_eq;
  3
  4use crate::db::billing_subscription::{StripeSubscriptionStatus, 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    // Test when an existing Zed Pro trial subscription is upgraded to Zed Pro.
 16    {
 17        let user_id = UserId(1);
 18
 19        let now = Utc::now();
 20
 21        let trial_period_start_at = now - Duration::days(14);
 22        let trial_period_end_at = now;
 23
 24        let new_period_start_at = now;
 25        let new_period_end_at = now + Duration::days(30);
 26
 27        let existing_subscription = billing_subscription::Model {
 28            kind: Some(SubscriptionKind::ZedProTrial),
 29            stripe_current_period_start: Some(trial_period_start_at.timestamp()),
 30            stripe_current_period_end: Some(trial_period_end_at.timestamp()),
 31            ..Default::default()
 32        };
 33
 34        let existing_usage = db
 35            .create_subscription_usage(
 36                user_id,
 37                trial_period_start_at,
 38                trial_period_end_at,
 39                SubscriptionKind::ZedProTrial,
 40                25,
 41                1_000,
 42            )
 43            .await
 44            .unwrap();
 45
 46        let transferred_usage = db
 47            .transfer_existing_subscription_usage(
 48                user_id,
 49                &existing_subscription,
 50                Some(SubscriptionKind::ZedPro),
 51                StripeSubscriptionStatus::Active,
 52                new_period_start_at,
 53                new_period_end_at,
 54            )
 55            .await
 56            .unwrap();
 57
 58        assert!(
 59            transferred_usage.is_some(),
 60            "subscription usage not transferred successfully"
 61        );
 62        let transferred_usage = transferred_usage.unwrap();
 63
 64        assert_eq!(
 65            transferred_usage.model_requests,
 66            existing_usage.model_requests
 67        );
 68        assert_eq!(
 69            transferred_usage.edit_predictions,
 70            existing_usage.edit_predictions
 71        );
 72    }
 73
 74    // Test when an existing Zed Pro trial subscription is canceled.
 75    {
 76        let user_id = UserId(2);
 77
 78        let now = Utc::now();
 79
 80        let trial_period_start_at = now - Duration::days(14);
 81        let trial_period_end_at = now;
 82
 83        let existing_subscription = billing_subscription::Model {
 84            kind: Some(SubscriptionKind::ZedProTrial),
 85            stripe_current_period_start: Some(trial_period_start_at.timestamp()),
 86            stripe_current_period_end: Some(trial_period_end_at.timestamp()),
 87            ..Default::default()
 88        };
 89
 90        let _existing_usage = db
 91            .create_subscription_usage(
 92                user_id,
 93                trial_period_start_at,
 94                trial_period_end_at,
 95                SubscriptionKind::ZedProTrial,
 96                25,
 97                1_000,
 98            )
 99            .await
100            .unwrap();
101
102        let transferred_usage = db
103            .transfer_existing_subscription_usage(
104                user_id,
105                &existing_subscription,
106                Some(SubscriptionKind::ZedPro),
107                StripeSubscriptionStatus::Canceled,
108                trial_period_start_at,
109                trial_period_end_at,
110            )
111            .await
112            .unwrap();
113
114        assert!(
115            transferred_usage.is_none(),
116            "subscription usage was transferred when it should not have been"
117        );
118    }
119}