collab: Factor out subscription kind determination (#30955)

Marshall Bowers created 6 months ago

This PR factors out the code that determines the `SubscriptionKind` into
a separate method for reusability purposes.

Release Notes:

- N/A

Change summary

crates/collab/src/api/billing.rs    | 30 +++++++-----------------------
crates/collab/src/stripe_billing.rs | 27 ++++++++++++++++++++++++++-
2 files changed, 33 insertions(+), 24 deletions(-)

Detailed changes

crates/collab/src/api/billing.rs

@@ -991,29 +991,13 @@ async fn handle_customer_subscription_event(
 
     log::info!("handling Stripe {} event: {}", event.type_, event.id);
 
-    let subscription_kind = maybe!(async {
-        let stripe_billing = app.stripe_billing.clone()?;
-
-        let zed_pro_price_id = stripe_billing.zed_pro_price_id().await.ok()?;
-        let zed_free_price_id = stripe_billing.zed_free_price_id().await.ok()?;
-
-        subscription.items.data.iter().find_map(|item| {
-            let price = item.price.as_ref()?;
-
-            if price.id == zed_pro_price_id {
-                Some(if subscription.status == SubscriptionStatus::Trialing {
-                    SubscriptionKind::ZedProTrial
-                } else {
-                    SubscriptionKind::ZedPro
-                })
-            } else if price.id == zed_free_price_id {
-                Some(SubscriptionKind::ZedFree)
-            } else {
-                None
-            }
-        })
-    })
-    .await;
+    let subscription_kind = if let Some(stripe_billing) = &app.stripe_billing {
+        stripe_billing
+            .determine_subscription_kind(&subscription)
+            .await
+    } else {
+        None
+    };
 
     let billing_customer =
         find_or_create_billing_customer(app, stripe_client, subscription.customer)

crates/collab/src/stripe_billing.rs

@@ -1,12 +1,13 @@
 use std::sync::Arc;
 
 use crate::Result;
+use crate::db::billing_subscription::SubscriptionKind;
 use crate::llm::AGENT_EXTENDED_TRIAL_FEATURE_FLAG;
 use anyhow::{Context as _, anyhow};
 use chrono::Utc;
 use collections::HashMap;
 use serde::{Deserialize, Serialize};
-use stripe::PriceId;
+use stripe::{PriceId, SubscriptionStatus};
 use tokio::sync::RwLock;
 use uuid::Uuid;
 
@@ -97,6 +98,30 @@ impl StripeBilling {
             .ok_or_else(|| crate::Error::Internal(anyhow!("no price found for {lookup_key:?}")))
     }
 
+    pub async fn determine_subscription_kind(
+        &self,
+        subscription: &stripe::Subscription,
+    ) -> Option<SubscriptionKind> {
+        let zed_pro_price_id = self.zed_pro_price_id().await.ok()?;
+        let zed_free_price_id = self.zed_free_price_id().await.ok()?;
+
+        subscription.items.data.iter().find_map(|item| {
+            let price = item.price.as_ref()?;
+
+            if price.id == zed_pro_price_id {
+                Some(if subscription.status == SubscriptionStatus::Trialing {
+                    SubscriptionKind::ZedProTrial
+                } else {
+                    SubscriptionKind::ZedPro
+                })
+            } else if price.id == zed_free_price_id {
+                Some(SubscriptionKind::ZedFree)
+            } else {
+                None
+            }
+        })
+    }
+
     pub async fn subscribe_to_price(
         &self,
         subscription_id: &stripe::SubscriptionId,