From 0159706d7543ba4501b1b1f28106953e4020fec0 Mon Sep 17 00:00:00 2001
From: dishenwang2023 <dishenwang1021@gmail.com>
Date: Mon, 27 May 2024 15:49:08 -0400
Subject: [PATCH] Fix stripe disabled functions for organization creation

---
 .../modules/organization/invoice/controller.py |  4 ++--
 .../organization/invoice/models/entities.py    | 11 +++++++++++
 .../enterprise/modules/organization/service.py | 18 ++++++++++++++----
 3 files changed, 27 insertions(+), 6 deletions(-)

diff --git a/services/enterprise/modules/organization/invoice/controller.py b/services/enterprise/modules/organization/invoice/controller.py
index 3cf89285..66255f21 100644
--- 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
index 0f9f56c7..2f14d64c 100644
--- 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
index db65a847..75853336 100644
--- 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,