import contextlib from dataclasses import dataclass from typing import Any, Optional import stripe from django.core.management.base import BaseCommand from django.utils.timezone import now as timezone_now from typing_extensions import override from corporate.lib.stripe import RealmBillingSession, add_months from corporate.models import Customer, CustomerPlan, LicenseLedger from zerver.actions.create_realm import do_create_realm from zerver.actions.create_user import do_create_user from zerver.actions.streams import bulk_add_subscriptions from zerver.apps import flush_cache from zerver.lib.streams import create_stream_if_needed from zerver.models import Realm, UserProfile, get_realm from zproject.config import get_secret @dataclass class CustomerProfile: unique_id: str billing_schedule: int = CustomerPlan.ANNUAL tier: Optional[int] = None automanage_licenses: bool = False status: int = CustomerPlan.ACTIVE sponsorship_pending: bool = False is_sponsored: bool = False card: str = "" class Command(BaseCommand): help = "Populate database with different types of realms that can exist." @override def handle(self, *args: Any, **options: Any) -> None: # Create a realm for each plan type customer_profiles = [ # NOTE: The unique_id has to be less than 40 characters. CustomerProfile(unique_id="sponsorship-pending", sponsorship_pending=True), CustomerProfile( unique_id="annual-free", billing_schedule=CustomerPlan.ANNUAL, ), CustomerProfile( unique_id="annual-standard", billing_schedule=CustomerPlan.ANNUAL, tier=CustomerPlan.STANDARD, ), CustomerProfile( unique_id="annual-plus", billing_schedule=CustomerPlan.ANNUAL, tier=CustomerPlan.PLUS, ), CustomerProfile( unique_id="monthly-free", billing_schedule=CustomerPlan.MONTHLY, ), CustomerProfile( unique_id="monthly-standard", billing_schedule=CustomerPlan.MONTHLY, tier=CustomerPlan.STANDARD, ), CustomerProfile( unique_id="monthly-plus", billing_schedule=CustomerPlan.MONTHLY, tier=CustomerPlan.PLUS, ), CustomerProfile( unique_id="downgrade-end-of-cycle", billing_schedule=CustomerPlan.MONTHLY, tier=CustomerPlan.STANDARD, status=CustomerPlan.DOWNGRADE_AT_END_OF_CYCLE, ), CustomerProfile( unique_id="standard-automanage-licenses", billing_schedule=CustomerPlan.MONTHLY, tier=CustomerPlan.STANDARD, automanage_licenses=True, ), CustomerProfile( unique_id="standard-automatic-card", billing_schedule=CustomerPlan.MONTHLY, tier=CustomerPlan.STANDARD, card="pm_card_visa", ), CustomerProfile( unique_id="sponsored", is_sponsored=True, billing_schedule=CustomerPlan.MONTHLY, tier=CustomerPlan.STANDARD, ), ] # Create a realm for each customer profile for customer_profile in customer_profiles: unique_id = customer_profile.unique_id if customer_profile.tier is None: plan_type = Realm.PLAN_TYPE_LIMITED elif customer_profile.tier == CustomerPlan.STANDARD and customer_profile.is_sponsored: plan_type = Realm.PLAN_TYPE_STANDARD_FREE elif customer_profile.tier == CustomerPlan.STANDARD: plan_type = Realm.PLAN_TYPE_STANDARD elif customer_profile.tier == CustomerPlan.PLUS: plan_type = Realm.PLAN_TYPE_PLUS else: raise AssertionError("Unexpected tier!") plan_name = Realm.ALL_PLAN_TYPES[plan_type] # Delete existing realm with this name with contextlib.suppress(Realm.DoesNotExist): get_realm(unique_id).delete() # Because we just deleted a bunch of objects in the database # directly (rather than deleting individual objects in Django, # in which case our post_save hooks would have flushed the # individual objects from memcached for us), we need to flush # memcached in order to ensure deleted objects aren't still # present in the memcached cache. flush_cache(None) realm = do_create_realm( string_id=unique_id, name=unique_id, description=unique_id, plan_type=plan_type, ) # Create a user with billing access full_name = f"{plan_name}-admin" email = f"{full_name}@zulip.com" user = do_create_user( email, full_name, realm, full_name, role=UserProfile.ROLE_REALM_OWNER, acting_user=None, ) stream, _ = create_stream_if_needed( realm, "all", ) bulk_add_subscriptions(realm, [stream], [user], acting_user=None) if customer_profile.sponsorship_pending: customer = Customer.objects.create( realm=realm, sponsorship_pending=customer_profile.sponsorship_pending, ) continue if customer_profile.tier is None: continue billing_session = RealmBillingSession(user) customer = billing_session.update_or_create_stripe_customer() assert customer.stripe_customer_id is not None # Add a test card to the customer. if customer_profile.card: # Set the Stripe API key stripe.api_key = get_secret("stripe_secret_key") # Create a card payment method and attach it to the customer payment_method = stripe.PaymentMethod.create( type="card", card={"token": "tok_visa"}, ) # Attach the payment method to the customer stripe.PaymentMethod.attach(payment_method.id, customer=customer.stripe_customer_id) # Set the default payment method for the customer stripe.Customer.modify( customer.stripe_customer_id, invoice_settings={"default_payment_method": payment_method.id}, ) months = 12 if customer_profile.billing_schedule == CustomerPlan.MONTHLY: months = 1 next_invoice_date = add_months(timezone_now(), months) customer_plan = CustomerPlan.objects.create( customer=customer, billing_cycle_anchor=timezone_now(), billing_schedule=customer_profile.billing_schedule, tier=customer_profile.tier, price_per_license=3, automanage_licenses=customer_profile.automanage_licenses, status=customer_profile.status, charge_automatically=True, next_invoice_date=next_invoice_date, ) LicenseLedger.objects.create( licenses=10, licenses_at_next_renewal=10, event_time=timezone_now(), is_renewal=True, plan=customer_plan, )