billing: Update license ledger when users are added and removed.

This commit is contained in:
Rishi Gupta
2019-01-25 17:36:37 -08:00
parent 03c71dad17
commit 7c11fe819a
4 changed files with 104 additions and 4 deletions

View File

@@ -17,7 +17,6 @@ from zerver.lib.exceptions import JsonableError
from zerver.lib.logging_util import log_to_file
from zerver.lib.timestamp import datetime_to_timestamp, timestamp_to_datetime
from zerver.lib.utils import generate_random_token
from zerver.lib.actions import do_change_plan_type
from zerver.models import Realm, UserProfile, RealmAuditLog
from corporate.models import Customer, CustomerPlan, LicenseLedger, \
get_active_plan
@@ -191,7 +190,7 @@ def do_replace_payment_source(user: UserProfile, stripe_token: str) -> stripe.Cu
# event_times in the past or future
# TODO handle downgrade
def add_plan_renewal_to_license_ledger_if_needed(plan: CustomerPlan, event_time: datetime) -> LicenseLedger:
last_ledger_entry = LicenseLedger.objects.filter(plan=plan).order_by('-event_time').first()
last_ledger_entry = LicenseLedger.objects.filter(plan=plan).order_by('-id').first()
plan_renewal_date = next_renewal_date(plan)
if plan_renewal_date < event_time:
if not LicenseLedger.objects.filter(
@@ -332,8 +331,30 @@ def process_initial_upgrade(user: UserProfile, licenses: int, automanage_license
statement_descriptor='Zulip Standard')
stripe.Invoice.finalize_invoice(stripe_invoice)
from zerver.lib.actions import do_change_plan_type
do_change_plan_type(realm, Realm.STANDARD)
def update_license_ledger_for_automanaged_plan(realm: Realm, plan: CustomerPlan,
event_time: datetime) -> None:
last_ledger_entry = add_plan_renewal_to_license_ledger_if_needed(plan, event_time)
# todo: handle downgrade, where licenses_at_next_renewal should be 0
licenses_at_next_renewal = get_seat_count(realm)
licenses = max(licenses_at_next_renewal, last_ledger_entry.licenses)
LicenseLedger.objects.create(
plan=plan, event_time=event_time, licenses=licenses,
licenses_at_next_renewal=licenses_at_next_renewal)
def update_license_ledger_if_needed(realm: Realm, event_time: datetime) -> None:
customer = Customer.objects.filter(realm=realm).first()
if customer is None:
return
plan = get_active_plan(customer)
if plan is None:
return
if not plan.automanage_licenses:
return
update_license_ledger_for_automanaged_plan(realm, plan, event_time)
def attach_discount_to_realm(user: UserProfile, discount: Decimal) -> None:
customer = Customer.objects.filter(realm=user.realm).first()
if customer is None:

View File

@@ -29,7 +29,8 @@ from corporate.lib.stripe import catch_stripe_errors, attach_discount_to_realm,
DEFAULT_INVOICE_DAYS_UNTIL_DUE, MIN_INVOICED_LICENSES, do_create_customer, \
add_months, next_month, next_renewal_date, renewal_amount, \
compute_plan_parameters, update_or_create_stripe_customer, \
process_initial_upgrade, add_plan_renewal_to_license_ledger_if_needed
process_initial_upgrade, add_plan_renewal_to_license_ledger_if_needed, \
update_license_ledger_if_needed, update_license_ledger_for_automanaged_plan
from corporate.models import Customer, CustomerPlan, LicenseLedger
from corporate.views import payment_method_string
import corporate.urls
@@ -964,3 +965,70 @@ class LicenseLedgerTest(ZulipTestCase):
# Plan needs to renew, but we already added the plan_renewal ledger entry
add_plan_renewal_to_license_ledger_if_needed(plan, self.next_year + timedelta(seconds=1))
self.assertEqual(LicenseLedger.objects.count(), 2)
def test_update_license_ledger_if_needed(self) -> None:
realm = get_realm('zulip')
# Test no Customer
update_license_ledger_if_needed(realm, self.now)
self.assertFalse(LicenseLedger.objects.exists())
# Test plan not automanaged
self.local_upgrade(self.seat_count + 1, False, CustomerPlan.ANNUAL, 'token')
self.assertEqual(LicenseLedger.objects.count(), 1)
update_license_ledger_if_needed(realm, self.now)
self.assertEqual(LicenseLedger.objects.count(), 1)
# Test no active plan
plan = CustomerPlan.objects.get()
plan.automanage_licenses = True
plan.status = CustomerPlan.ENDED
plan.save(update_fields=['automanage_licenses', 'status'])
update_license_ledger_if_needed(realm, self.now)
self.assertEqual(LicenseLedger.objects.count(), 1)
# Test update needed
plan.status = CustomerPlan.ACTIVE
plan.save(update_fields=['status'])
update_license_ledger_if_needed(realm, self.now)
self.assertEqual(LicenseLedger.objects.count(), 2)
def test_update_license_ledger_for_automanaged_plan(self) -> None:
realm = get_realm('zulip')
with patch('corporate.lib.stripe.timezone_now', return_value=self.now):
self.local_upgrade(self.seat_count, True, CustomerPlan.ANNUAL, 'token')
plan = CustomerPlan.objects.first()
# Simple increase
with patch('corporate.lib.stripe.get_seat_count', return_value=23):
update_license_ledger_for_automanaged_plan(realm, plan, self.now)
# Decrease
with patch('corporate.lib.stripe.get_seat_count', return_value=20):
update_license_ledger_for_automanaged_plan(realm, plan, self.now)
# Increase, but not past high watermark
with patch('corporate.lib.stripe.get_seat_count', return_value=21):
update_license_ledger_for_automanaged_plan(realm, plan, self.now)
# Increase, but after renewal date, and below last year's high watermark
with patch('corporate.lib.stripe.get_seat_count', return_value=22):
update_license_ledger_for_automanaged_plan(realm, plan, self.next_year + timedelta(seconds=1))
ledger_entries = list(LicenseLedger.objects.values_list(
'is_renewal', 'event_time', 'licenses', 'licenses_at_next_renewal').order_by('id'))
self.assertEqual(ledger_entries,
[(True, self.now, self.seat_count, self.seat_count),
(False, self.now, 23, 23),
(False, self.now, 23, 20),
(False, self.now, 23, 21),
(True, self.next_year, 21, 21),
(False, self.next_year + timedelta(seconds=1), 22, 22)])
def test_user_changes(self) -> None:
self.local_upgrade(self.seat_count, True, CustomerPlan.ANNUAL, 'token')
user = do_create_user('email', 'password', get_realm('zulip'), 'name', 'name')
do_deactivate_user(user)
do_reactivate_user(user)
# Not a proper use of do_activate_user, but fine for this test
do_activate_user(user)
ledger_entries = list(LicenseLedger.objects.values_list(
'is_renewal', 'licenses', 'licenses_at_next_renewal').order_by('id'))
self.assertEqual(ledger_entries,
[(True, self.seat_count, self.seat_count),
(False, self.seat_count + 1, self.seat_count + 1),
(False, self.seat_count + 1, self.seat_count),
(False, self.seat_count + 1, self.seat_count + 1),
(False, self.seat_count + 1, self.seat_count + 1)])

View File

@@ -149,6 +149,9 @@ from zerver.lib.types import ProfileFieldData
from analytics.models import StreamCount
if settings.BILLING_ENABLED:
from corporate.lib.stripe import update_license_ledger_if_needed
import ujson
import time
import re
@@ -488,6 +491,8 @@ def do_create_user(email: str, password: Optional[str], realm: Realm, full_name:
event_type=RealmAuditLog.USER_CREATED, event_time=event_time)
do_increment_logging_stat(user_profile.realm, COUNT_STATS['active_users_log:is_bot:day'],
user_profile.is_bot, event_time)
if settings.BILLING_ENABLED:
update_license_ledger_if_needed(user_profile.realm, event_time)
notify_created_user(user_profile)
if bot_type:
@@ -513,6 +518,8 @@ def do_activate_user(user_profile: UserProfile) -> None:
event_type=RealmAuditLog.USER_ACTIVATED, event_time=event_time)
do_increment_logging_stat(user_profile.realm, COUNT_STATS['active_users_log:is_bot:day'],
user_profile.is_bot, event_time)
if settings.BILLING_ENABLED:
update_license_ledger_if_needed(user_profile.realm, event_time)
notify_created_user(user_profile)
@@ -528,6 +535,8 @@ def do_reactivate_user(user_profile: UserProfile, acting_user: Optional[UserProf
acting_user=acting_user)
do_increment_logging_stat(user_profile.realm, COUNT_STATS['active_users_log:is_bot:day'],
user_profile.is_bot, event_time)
if settings.BILLING_ENABLED:
update_license_ledger_if_needed(user_profile.realm, event_time)
notify_created_user(user_profile)
@@ -704,6 +713,8 @@ def do_deactivate_user(user_profile: UserProfile,
event_type=RealmAuditLog.USER_DEACTIVATED, event_time=event_time)
do_increment_logging_stat(user_profile.realm, COUNT_STATS['active_users_log:is_bot:day'],
user_profile.is_bot, event_time, increment=-1)
if settings.BILLING_ENABLED:
update_license_ledger_if_needed(user_profile.realm, event_time)
event = dict(type="realm_user", op="remove",
person=dict(email=user_profile.email,

View File

@@ -464,7 +464,7 @@ class LoginTest(ZulipTestCase):
with queries_captured() as queries:
self.register(self.nonreg_email('test'), "test")
# Ensure the number of queries we make is not O(streams)
self.assert_length(queries, 73)
self.assert_length(queries, 74)
user_profile = self.nonreg_user('test')
self.assertEqual(get_session_dict_user(self.client.session), user_profile.id)
self.assertFalse(user_profile.enable_stream_desktop_notifications)