Yehowshua/ZIm
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

collab: Don't try to transfer usage when a Zed Pro trial is canceled (#29843)

Browse source 
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
This commit is contained in:
Marshall Bowers 2025-05-03 10:57:54 -04:00 • committed by GitHub
parent 7f8e3fd482
commit 12c26a4fa6
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 102 additions and 50 deletions
Download patch file Download diff file Expand all files Collapse all files

1
crates/collab/src/api/billing.rs
View file

@ -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,
)

7
crates/collab/src/llm/db/queries/subscription_usages.rs
View file

@ -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"))?;

144
crates/collab/src/llm/db/tests/subscription_usage_tests.rs
View file

@ -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"
);
}
}