mirror of https://github.com/zulip/zulip.git
populate_billing_realms: Add script to populate different realms.
These realms are based on different cases of realms which can exist while billing them and hence are useful for testing.
This commit is contained in:
Aman Agrawal
Tim Abbott
parent
0b3f7a5a6b
commit
9ebb3f06b3
|
|
@ -655,6 +655,15 @@ class Realm(models.Model): # type: ignore[django-manager-missing] # django-stub
|
|||
PLAN_TYPE_STANDARD = 3
|
||||
PLAN_TYPE_STANDARD_FREE = 4
|
||||
PLAN_TYPE_PLUS = 10
|
||||
|
||||
# Used for creating realms with different plan types.
|
||||
ALL_PLAN_TYPES = {
|
||||
PLAN_TYPE_SELF_HOSTED: "self-hosted-plan",
|
||||
PLAN_TYPE_LIMITED: "limited-plan",
|
||||
PLAN_TYPE_STANDARD: "standard-plan",
|
||||
PLAN_TYPE_STANDARD_FREE: "standard-free-plan",
|
||||
PLAN_TYPE_PLUS: "plus-plan",
|
||||
}
|
||||
plan_type = models.PositiveSmallIntegerField(default=PLAN_TYPE_SELF_HOSTED)
|
||||
|
||||
# This value is also being used in web/src/settings_bots.bot_creation_policy_values.
|
||||
|
|
|
|||
|
|
@ -0,0 +1,190 @@
|
|||
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 update_or_create_stripe_customer
|
||||
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-standard",
|
||||
billing_schedule=CustomerPlan.ANNUAL,
|
||||
tier=CustomerPlan.STANDARD,
|
||||
),
|
||||
CustomerProfile(
|
||||
unique_id="annual-plus",
|
||||
billing_schedule=CustomerPlan.ANNUAL,
|
||||
tier=CustomerPlan.PLUS,
|
||||
),
|
||||
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
|
||||
|
||||
assert customer_profile.tier is not None
|
||||
|
||||
customer = update_or_create_stripe_customer(user)
|
||||
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},
|
||||
)
|
||||
|
||||
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,
|
||||
)
|
||||
|
||||
LicenseLedger.objects.create(
|
||||
licenses=10,
|
||||
licenses_at_next_renewal=10,
|
||||
event_time=timezone_now(),
|
||||
is_renewal=True,
|
||||
plan=customer_plan,
|
||||
)
|
||||
Loading…
Reference in New Issue