From 12c26a4fa69c1d264a134b7337e16b75cceaa4d0 Mon Sep 17 00:00:00 2001 From: Marshall Bowers Date: Sat, 3 May 2025 10:57:54 -0400 Subject: [PATCH] collab: Don't try to transfer usage when a Zed Pro trial is canceled (#29843) 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/A --- crates/collab/src/api/billing.rs | 1 + .../src/llm/db/queries/subscription_usages.rs | 7 +- .../llm/db/tests/subscription_usage_tests.rs | 144 ++++++++++++------ 3 files changed, 102 insertions(+), 50 deletions(-) diff --git a/crates/collab/src/api/billing.rs b/crates/collab/src/api/billing.rs index 3824e834c3..b8ac82a09b 100644 --- a/crates/collab/src/api/billing.rs +++ b/crates/collab/src/api/billing.rs @@ -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, ) diff --git a/crates/collab/src/llm/db/queries/subscription_usages.rs b/crates/collab/src/llm/db/queries/subscription_usages.rs index 162bbd274d..48a3d54feb 100644 --- a/crates/collab/src/llm/db/queries/subscription_usages.rs +++ b/crates/collab/src/llm/db/queries/subscription_usages.rs @@ -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, + new_subscription_status: StripeSubscriptionStatus, new_period_start_at: DateTimeUtc, new_period_end_at: DateTimeUtc, ) -> Result> { 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"))?; diff --git a/crates/collab/src/llm/db/tests/subscription_usage_tests.rs b/crates/collab/src/llm/db/tests/subscription_usage_tests.rs index 84c92733ac..5ae49515be 100644 --- a/crates/collab/src/llm/db/tests/subscription_usage_tests.rs +++ b/crates/collab/src/llm/db/tests/subscription_usage_tests.rs @@ -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); + // Test when an existing Zed Pro trial subscription is upgraded to Zed Pro. + { + let user_id = UserId(1); - let now = Utc::now(); + let now = Utc::now(); - let trial_period_start_at = now - Duration::days(14); - let trial_period_end_at = 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 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_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 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(); + 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!( + 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 - ); + 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" + ); + } }