Fix stripe disabled functions for organization creation (#502)

Dataherald · May 27, 2024 · d1617e3 · d1617e3
1 parent e3ef091
commit d1617e3
Show file tree

Hide file tree

Showing 3 changed files with 27 additions and 6 deletions.
diff --git a/services/enterprise/modules/organization/invoice/controller.py b/services/enterprise/modules/organization/invoice/controller.py
@@ -19,7 +19,7 @@
 from utils.auth import Authorize, User, authenticate_user
 
 
-def check_stripe_disabled(request: Request):
+def is_stripe_disabled(request: Request):
     if invoice_settings.stripe_disabled:
         raise StripeDisabledError()
     return request
@@ -28,7 +28,7 @@ def check_stripe_disabled(request: Request):
 router = APIRouter(
     prefix="/organizations",
     responses={404: {"description": "Not found"}},
-    dependencies=[Depends(check_stripe_disabled)],
+    dependencies=[Depends(is_stripe_disabled)],
 )
 
 authorize = Authorize()

diff --git a/services/enterprise/modules/organization/invoice/models/entities.py b/services/enterprise/modules/organization/invoice/models/entities.py
@@ -70,3 +70,14 @@ class UsageInvoice(BaseModel):
     sql_generation_cost: int = 0
     finetuning_gpt_35_cost: int = 0
     finetuning_gpt_4_cost: int = 0
+
+
+class MockStripeCustomer(BaseModel):
+    id: str | None = None
+    name: str | None = None
+
+
+class MockStripeSubscription(BaseModel):
+    id: str | None = None
+    status: str | None = None
+    billing_cycle_anchor: int | None = None
diff --git a/services/enterprise/modules/organization/service.py b/services/enterprise/modules/organization/service.py
@@ -4,6 +4,8 @@
 from modules.organization.invoice.models.entities import (
     Credit,
     InvoiceDetails,
+    MockStripeCustomer,
+    MockStripeSubscription,
     PaymentPlan,
     RecordStatus,
 )
@@ -62,8 +64,12 @@ def add_organization(
             )
         organization = Organization(**org_request.dict())
 
-        customer = self.billing.create_customer(organization.name)
-        subscription = self.billing.create_subscription(customer.id)
+        if invoice_settings.stripe_disabled:
+            customer = MockStripeCustomer()
+            subscription = MockStripeSubscription()
+        else:
+            customer = self.billing.create_customer(organization.name)
+            subscription = self.billing.create_subscription(customer.id)
         # default organization plan is CREDIT_ONLY
         organization.invoice_details = InvoiceDetails(
             plan=PaymentPlan.CREDIT_ONLY,
@@ -161,8 +167,12 @@ def add_organization_by_slack_installation(
             owner=slack_installation_request.user.id,
         )
 
-        customer = self.billing.create_customer(organization.name)
-        subscription = self.billing.create_subscription(customer.id)
+        if invoice_settings.stripe_disabled:
+            customer = MockStripeCustomer()
+            subscription = MockStripeSubscription()
+        else:
+            customer = self.billing.create_customer(organization.name)
+            subscription = self.billing.create_subscription(customer.id)
         organization.invoice_details = InvoiceDetails(
             plan=PaymentPlan.CREDIT_ONLY,
             stripe_customer_id=customer.id,