Marshall Bowers created
This PR fixes an issue where we would erroneously try to transfer
existing subscription usage when a Zed Pro trial was canceled.
Release Notes:
- N/Acrates/collab/src/api/billing.rs | 1
crates/collab/src/llm/db/queries/subscription_usages.rs | 7
crates/collab/src/llm/db/tests/subscription_usage_tests.rs | 160 +++++--
3 files changed, 110 insertions(+), 58 deletions(-)
@@ -1040,6 +1040,7 @@ async fn handle_customer_subscription_event(
billing_customer.user_id,
&existing_subscription,
subscription_kind,
+ subscription.status.into(),
new_period_start_at,
new_period_end_at,
)
@@ -1,7 +1,7 @@
use chrono::Timelike;
use time::PrimitiveDateTime;
-use crate::db::billing_subscription::SubscriptionKind;
+use crate::db::billing_subscription::{StripeSubscriptionStatus, SubscriptionKind};
use crate::db::{UserId, billing_subscription};
use super::*;
@@ -120,12 +120,13 @@ impl LlmDatabase {
user_id: UserId,
existing_subscription: &billing_subscription::Model,
new_subscription_kind: Option<SubscriptionKind>,
+ new_subscription_status: StripeSubscriptionStatus,
new_period_start_at: DateTimeUtc,
new_period_end_at: DateTimeUtc,
) -> Result<Option<subscription_usage::Model>> {
self.transaction(|tx| async move {
- match existing_subscription.kind {
- Some(SubscriptionKind::ZedProTrial) => {
+ match (existing_subscription.kind, new_subscription_status) {
+ (Some(SubscriptionKind::ZedProTrial), StripeSubscriptionStatus::Active) => {
let trial_period_start_at = existing_subscription
.current_period_start_at()
.ok_or_else(|| anyhow!("No trial subscription period start"))?;
@@ -1,7 +1,7 @@
use chrono::{Duration, Utc};
use pretty_assertions::assert_eq;
-use crate::db::billing_subscription::SubscriptionKind;
+use crate::db::billing_subscription::{StripeSubscriptionStatus, SubscriptionKind};
use crate::db::{UserId, billing_subscription};
use crate::llm::db::LlmDatabase;
use crate::test_llm_db;
@@ -12,58 +12,108 @@ test_llm_db!(
);
async fn test_transfer_existing_subscription_usage(db: &mut LlmDatabase) {
- let user_id = UserId(1);
-
- let now = Utc::now();
-
- let trial_period_start_at = now - Duration::days(14);
- let trial_period_end_at = now;
-
- let new_period_start_at = now;
- let new_period_end_at = now + Duration::days(30);
-
- let existing_subscription = billing_subscription::Model {
- kind: Some(SubscriptionKind::ZedProTrial),
- stripe_current_period_start: Some(trial_period_start_at.timestamp()),
- stripe_current_period_end: Some(trial_period_end_at.timestamp()),
- ..Default::default()
- };
-
- let existing_usage = db
- .create_subscription_usage(
- user_id,
- trial_period_start_at,
- trial_period_end_at,
- SubscriptionKind::ZedProTrial,
- 25,
- 1_000,
- )
- .await
- .unwrap();
-
- let transferred_usage = db
- .transfer_existing_subscription_usage(
- user_id,
- &existing_subscription,
- Some(SubscriptionKind::ZedPro),
- new_period_start_at,
- new_period_end_at,
- )
- .await
- .unwrap();
-
- assert!(
- transferred_usage.is_some(),
- "subscription usage not transferred successfully"
- );
- let transferred_usage = transferred_usage.unwrap();
-
- assert_eq!(
- transferred_usage.model_requests,
- existing_usage.model_requests
- );
- assert_eq!(
- transferred_usage.edit_predictions,
- existing_usage.edit_predictions
- );
+ // Test when an existing Zed Pro trial subscription is upgraded to Zed Pro.
+ {
+ let user_id = UserId(1);
+
+ let now = Utc::now();
+
+ let trial_period_start_at = now - Duration::days(14);
+ let trial_period_end_at = now;
+
+ let new_period_start_at = now;
+ let new_period_end_at = now + Duration::days(30);
+
+ let existing_subscription = billing_subscription::Model {
+ kind: Some(SubscriptionKind::ZedProTrial),
+ stripe_current_period_start: Some(trial_period_start_at.timestamp()),
+ stripe_current_period_end: Some(trial_period_end_at.timestamp()),
+ ..Default::default()
+ };
+
+ let existing_usage = db
+ .create_subscription_usage(
+ user_id,
+ trial_period_start_at,
+ trial_period_end_at,
+ SubscriptionKind::ZedProTrial,
+ 25,
+ 1_000,
+ )
+ .await
+ .unwrap();
+
+ let transferred_usage = db
+ .transfer_existing_subscription_usage(
+ user_id,
+ &existing_subscription,
+ Some(SubscriptionKind::ZedPro),
+ StripeSubscriptionStatus::Active,
+ new_period_start_at,
+ new_period_end_at,
+ )
+ .await
+ .unwrap();
+
+ assert!(
+ transferred_usage.is_some(),
+ "subscription usage not transferred successfully"
+ );
+ let transferred_usage = transferred_usage.unwrap();
+
+ assert_eq!(
+ transferred_usage.model_requests,
+ existing_usage.model_requests
+ );
+ assert_eq!(
+ transferred_usage.edit_predictions,
+ existing_usage.edit_predictions
+ );
+ }
+
+ // Test when an existing Zed Pro trial subscription is canceled.
+ {
+ let user_id = UserId(2);
+
+ let now = Utc::now();
+
+ let trial_period_start_at = now - Duration::days(14);
+ let trial_period_end_at = now;
+
+ let existing_subscription = billing_subscription::Model {
+ kind: Some(SubscriptionKind::ZedProTrial),
+ stripe_current_period_start: Some(trial_period_start_at.timestamp()),
+ stripe_current_period_end: Some(trial_period_end_at.timestamp()),
+ ..Default::default()
+ };
+
+ let _existing_usage = db
+ .create_subscription_usage(
+ user_id,
+ trial_period_start_at,
+ trial_period_end_at,
+ SubscriptionKind::ZedProTrial,
+ 25,
+ 1_000,
+ )
+ .await
+ .unwrap();
+
+ let transferred_usage = db
+ .transfer_existing_subscription_usage(
+ user_id,
+ &existing_subscription,
+ Some(SubscriptionKind::ZedPro),
+ StripeSubscriptionStatus::Canceled,
+ trial_period_start_at,
+ trial_period_end_at,
+ )
+ .await
+ .unwrap();
+
+ assert!(
+ transferred_usage.is_none(),
+ "subscription usage was transferred when it should not have been"
+ );
+ }
}