billing: Rename get_seat_count to get_latest_seat_count.

This will help as our billing system becomes more async to accommodate
on-prem billing.
This commit is contained in:
Rishi Gupta 2019-10-07 10:21:29 -07:00 committed by Tim Abbott
parent bea9e41fbd
commit 0b39263ec0
3 changed files with 34 additions and 34 deletions

View File

@ -38,7 +38,7 @@ CallableT = TypeVar('CallableT', bound=Callable[..., Any])
MIN_INVOICED_LICENSES = 30 MIN_INVOICED_LICENSES = 30
DEFAULT_INVOICE_DAYS_UNTIL_DUE = 30 DEFAULT_INVOICE_DAYS_UNTIL_DUE = 30
def get_seat_count(realm: Realm) -> int: def get_latest_seat_count(realm: Realm) -> int:
non_guests = UserProfile.objects.filter( non_guests = UserProfile.objects.filter(
realm=realm, is_active=True, is_bot=False).exclude(role=UserProfile.ROLE_GUEST).count() realm=realm, is_active=True, is_bot=False).exclude(role=UserProfile.ROLE_GUEST).count()
guests = UserProfile.objects.filter( guests = UserProfile.objects.filter(
@ -314,7 +314,7 @@ def process_initial_upgrade(user: UserProfile, licenses: int, automanage_license
with transaction.atomic(): with transaction.atomic():
# billed_licenses can greater than licenses if users are added between the start of # billed_licenses can greater than licenses if users are added between the start of
# this function (process_initial_upgrade) and now # this function (process_initial_upgrade) and now
billed_licenses = max(get_seat_count(realm), licenses) billed_licenses = max(get_latest_seat_count(realm), licenses)
plan_params = { plan_params = {
'automanage_licenses': automanage_licenses, 'automanage_licenses': automanage_licenses,
'charge_automatically': charge_automatically, 'charge_automatically': charge_automatically,
@ -371,7 +371,7 @@ def update_license_ledger_for_automanaged_plan(realm: Realm, plan: CustomerPlan,
last_ledger_entry = make_end_of_cycle_updates_if_needed(plan, event_time) last_ledger_entry = make_end_of_cycle_updates_if_needed(plan, event_time)
if last_ledger_entry is None: if last_ledger_entry is None:
return return
licenses_at_next_renewal = get_seat_count(realm) licenses_at_next_renewal = get_latest_seat_count(realm)
licenses = max(licenses_at_next_renewal, last_ledger_entry.licenses) licenses = max(licenses_at_next_renewal, last_ledger_entry.licenses)
LicenseLedger.objects.create( LicenseLedger.objects.create(
plan=plan, event_time=event_time, licenses=licenses, plan=plan, event_time=event_time, licenses=licenses,

View File

@ -23,7 +23,7 @@ from zerver.lib.test_classes import ZulipTestCase
from zerver.lib.timestamp import timestamp_to_datetime, datetime_to_timestamp from zerver.lib.timestamp import timestamp_to_datetime, datetime_to_timestamp
from zerver.models import Realm, UserProfile, get_realm, RealmAuditLog from zerver.models import Realm, UserProfile, get_realm, RealmAuditLog
from corporate.lib.stripe import catch_stripe_errors, attach_discount_to_realm, \ from corporate.lib.stripe import catch_stripe_errors, attach_discount_to_realm, \
get_seat_count, sign_string, unsign_string, \ get_latest_seat_count, sign_string, unsign_string, \
BillingError, StripeCardError, stripe_get_customer, \ BillingError, StripeCardError, stripe_get_customer, \
MIN_INVOICED_LICENSES, \ MIN_INVOICED_LICENSES, \
add_months, next_month, \ add_months, next_month, \
@ -215,9 +215,9 @@ class StripeTestCase(ZulipTestCase):
def setUp(self, *mocks: Mock) -> None: def setUp(self, *mocks: Mock) -> None:
super().setUp() super().setUp()
# This test suite is not robust to users being added in populate_db. The following # This test suite is not robust to users being added in populate_db. The following
# hack ensures get_seat_count is fixed, even as populate_db changes. # hack ensures get_latest_seat_count is fixed, even as populate_db changes.
realm = get_realm('zulip') realm = get_realm('zulip')
seat_count = get_seat_count(realm) seat_count = get_latest_seat_count(realm)
assert(seat_count >= 6) assert(seat_count >= 6)
for user in UserProfile.objects.filter(realm=realm, is_active=True, is_bot=False) \ for user in UserProfile.objects.filter(realm=realm, is_active=True, is_bot=False) \
.exclude(role=UserProfile.ROLE_GUEST).exclude(email__in=[ .exclude(role=UserProfile.ROLE_GUEST).exclude(email__in=[
@ -225,7 +225,7 @@ class StripeTestCase(ZulipTestCase):
self.example_email('iago')])[:seat_count-6]: self.example_email('iago')])[:seat_count-6]:
user.is_active = False user.is_active = False
user.save(update_fields=['is_active']) user.save(update_fields=['is_active'])
self.assertEqual(get_seat_count(realm), 6) self.assertEqual(get_latest_seat_count(realm), 6)
self.seat_count = 6 self.seat_count = 6
self.signed_seat_count, self.salt = sign_string(str(self.seat_count)) self.signed_seat_count, self.salt = sign_string(str(self.seat_count))
# Choosing dates with corresponding timestamps below 1500000000 so that they are # Choosing dates with corresponding timestamps below 1500000000 so that they are
@ -544,7 +544,7 @@ class StripeTest(StripeTestCase):
self.login(self.example_email("hamlet")) self.login(self.example_email("hamlet"))
new_seat_count = 23 new_seat_count = 23
# Change the seat count while the user is going through the upgrade flow # Change the seat count while the user is going through the upgrade flow
with patch('corporate.lib.stripe.get_seat_count', return_value=new_seat_count): with patch('corporate.lib.stripe.get_latest_seat_count', return_value=new_seat_count):
self.upgrade() self.upgrade()
stripe_customer_id = Customer.objects.first().stripe_customer_id stripe_customer_id = Customer.objects.first().stripe_customer_id
# Check that the Charge used the old quantity, not new_seat_count # Check that the Charge used the old quantity, not new_seat_count
@ -592,8 +592,8 @@ class StripeTest(StripeTestCase):
self.assertEqual('/upgrade/', response.url) self.assertEqual('/upgrade/', response.url)
# Try again, with a valid card, after they added a few users # Try again, with a valid card, after they added a few users
with patch('corporate.lib.stripe.get_seat_count', return_value=23): with patch('corporate.lib.stripe.get_latest_seat_count', return_value=23):
with patch('corporate.views.get_seat_count', return_value=23): with patch('corporate.views.get_latest_seat_count', return_value=23):
self.upgrade() self.upgrade()
customer = Customer.objects.get(realm=get_realm('zulip')) customer = Customer.objects.get(realm=get_realm('zulip'))
# It's impossible to create two Customers, but check that we didn't # It's impossible to create two Customers, but check that we didn't
@ -723,37 +723,37 @@ class StripeTest(StripeTestCase):
self.assertEqual(response.status_code, 302) self.assertEqual(response.status_code, 302)
self.assertEqual('/upgrade/', response.url) self.assertEqual('/upgrade/', response.url)
def test_get_seat_count(self) -> None: def test_get_latest_seat_count(self) -> None:
realm = get_realm("zulip") realm = get_realm("zulip")
initial_count = get_seat_count(realm) initial_count = get_latest_seat_count(realm)
user1 = UserProfile.objects.create(realm=realm, email='user1@zulip.com', pointer=-1) user1 = UserProfile.objects.create(realm=realm, email='user1@zulip.com', pointer=-1)
user2 = UserProfile.objects.create(realm=realm, email='user2@zulip.com', pointer=-1) user2 = UserProfile.objects.create(realm=realm, email='user2@zulip.com', pointer=-1)
self.assertEqual(get_seat_count(realm), initial_count + 2) self.assertEqual(get_latest_seat_count(realm), initial_count + 2)
# Test that bots aren't counted # Test that bots aren't counted
user1.is_bot = True user1.is_bot = True
user1.save(update_fields=['is_bot']) user1.save(update_fields=['is_bot'])
self.assertEqual(get_seat_count(realm), initial_count + 1) self.assertEqual(get_latest_seat_count(realm), initial_count + 1)
# Test that inactive users aren't counted # Test that inactive users aren't counted
do_deactivate_user(user2) do_deactivate_user(user2)
self.assertEqual(get_seat_count(realm), initial_count) self.assertEqual(get_latest_seat_count(realm), initial_count)
# Test guests # Test guests
# Adding a guest to a realm with a lot of members shouldn't change anything # Adding a guest to a realm with a lot of members shouldn't change anything
UserProfile.objects.create(realm=realm, email='user3@zulip.com', pointer=-1, role=UserProfile.ROLE_GUEST) UserProfile.objects.create(realm=realm, email='user3@zulip.com', pointer=-1, role=UserProfile.ROLE_GUEST)
self.assertEqual(get_seat_count(realm), initial_count) self.assertEqual(get_latest_seat_count(realm), initial_count)
# Test 1 member and 5 guests # Test 1 member and 5 guests
realm = Realm.objects.create(string_id='second', name='second') realm = Realm.objects.create(string_id='second', name='second')
UserProfile.objects.create(realm=realm, email='member@second.com', pointer=-1) UserProfile.objects.create(realm=realm, email='member@second.com', pointer=-1)
for i in range(5): for i in range(5):
UserProfile.objects.create(realm=realm, email='guest{}@second.com'.format(i), UserProfile.objects.create(realm=realm, email='guest{}@second.com'.format(i),
pointer=-1, role=UserProfile.ROLE_GUEST) pointer=-1, role=UserProfile.ROLE_GUEST)
self.assertEqual(get_seat_count(realm), 1) self.assertEqual(get_latest_seat_count(realm), 1)
# Test 1 member and 6 guests # Test 1 member and 6 guests
UserProfile.objects.create(realm=realm, email='guest5@second.com', pointer=-1, UserProfile.objects.create(realm=realm, email='guest5@second.com', pointer=-1,
role=UserProfile.ROLE_GUEST) role=UserProfile.ROLE_GUEST)
self.assertEqual(get_seat_count(realm), 2) self.assertEqual(get_latest_seat_count(realm), 2)
def test_sign_string(self) -> None: def test_sign_string(self) -> None:
string = "abc" string = "abc"
@ -892,7 +892,7 @@ class StripeTest(StripeTestCase):
# Verify that we still write LicenseLedger rows during the remaining # Verify that we still write LicenseLedger rows during the remaining
# part of the cycle # part of the cycle
with patch("corporate.lib.stripe.get_seat_count", return_value=20): with patch("corporate.lib.stripe.get_latest_seat_count", return_value=20):
update_license_ledger_if_needed(user.realm, self.now) update_license_ledger_if_needed(user.realm, self.now)
self.assertEqual(LicenseLedger.objects.order_by('-id').values_list( self.assertEqual(LicenseLedger.objects.order_by('-id').values_list(
'licenses', 'licenses_at_next_renewal').first(), (20, 20)) 'licenses', 'licenses_at_next_renewal').first(), (20, 20))
@ -907,7 +907,7 @@ class StripeTest(StripeTestCase):
mocked.reset_mock() mocked.reset_mock()
# Check that we downgrade properly if the cycle is over # Check that we downgrade properly if the cycle is over
with patch("corporate.lib.stripe.get_seat_count", return_value=30): with patch("corporate.lib.stripe.get_latest_seat_count", return_value=30):
update_license_ledger_if_needed(user.realm, self.next_year) update_license_ledger_if_needed(user.realm, self.next_year)
self.assertEqual(get_realm('zulip').plan_type, Realm.LIMITED) self.assertEqual(get_realm('zulip').plan_type, Realm.LIMITED)
self.assertEqual(CustomerPlan.objects.first().status, CustomerPlan.ENDED) self.assertEqual(CustomerPlan.objects.first().status, CustomerPlan.ENDED)
@ -915,7 +915,7 @@ class StripeTest(StripeTestCase):
'licenses', 'licenses_at_next_renewal').first(), (20, 20)) 'licenses', 'licenses_at_next_renewal').first(), (20, 20))
# Verify that we don't write LicenseLedger rows once we've downgraded # Verify that we don't write LicenseLedger rows once we've downgraded
with patch("corporate.lib.stripe.get_seat_count", return_value=40): with patch("corporate.lib.stripe.get_latest_seat_count", return_value=40):
update_license_ledger_if_needed(user.realm, self.next_year) update_license_ledger_if_needed(user.realm, self.next_year)
self.assertEqual(LicenseLedger.objects.order_by('-id').values_list( self.assertEqual(LicenseLedger.objects.order_by('-id').values_list(
'licenses', 'licenses_at_next_renewal').first(), (20, 20)) 'licenses', 'licenses_at_next_renewal').first(), (20, 20))
@ -931,7 +931,7 @@ class StripeTest(StripeTestCase):
self.assertIsNone(CustomerPlan.objects.first().next_invoice_date) self.assertIsNone(CustomerPlan.objects.first().next_invoice_date)
# Check that we don't call invoice_plan after that final call # Check that we don't call invoice_plan after that final call
with patch("corporate.lib.stripe.get_seat_count", return_value=50): with patch("corporate.lib.stripe.get_latest_seat_count", return_value=50):
update_license_ledger_if_needed(user.realm, self.next_year + timedelta(days=80)) update_license_ledger_if_needed(user.realm, self.next_year + timedelta(days=80))
with patch("corporate.lib.stripe.invoice_plan") as mocked: with patch("corporate.lib.stripe.invoice_plan") as mocked:
invoice_plans_as_needed(self.next_year + timedelta(days=400)) invoice_plans_as_needed(self.next_year + timedelta(days=400))
@ -1141,16 +1141,16 @@ class LicenseLedgerTest(StripeTestCase):
self.local_upgrade(self.seat_count, True, CustomerPlan.ANNUAL, 'token') self.local_upgrade(self.seat_count, True, CustomerPlan.ANNUAL, 'token')
plan = CustomerPlan.objects.first() plan = CustomerPlan.objects.first()
# Simple increase # Simple increase
with patch('corporate.lib.stripe.get_seat_count', return_value=23): with patch('corporate.lib.stripe.get_latest_seat_count', return_value=23):
update_license_ledger_for_automanaged_plan(realm, plan, self.now) update_license_ledger_for_automanaged_plan(realm, plan, self.now)
# Decrease # Decrease
with patch('corporate.lib.stripe.get_seat_count', return_value=20): with patch('corporate.lib.stripe.get_latest_seat_count', return_value=20):
update_license_ledger_for_automanaged_plan(realm, plan, self.now) update_license_ledger_for_automanaged_plan(realm, plan, self.now)
# Increase, but not past high watermark # Increase, but not past high watermark
with patch('corporate.lib.stripe.get_seat_count', return_value=21): with patch('corporate.lib.stripe.get_latest_seat_count', return_value=21):
update_license_ledger_for_automanaged_plan(realm, plan, self.now) update_license_ledger_for_automanaged_plan(realm, plan, self.now)
# Increase, but after renewal date, and below last year's high watermark # Increase, but after renewal date, and below last year's high watermark
with patch('corporate.lib.stripe.get_seat_count', return_value=22): with patch('corporate.lib.stripe.get_latest_seat_count', return_value=22):
update_license_ledger_for_automanaged_plan(realm, plan, self.next_year + timedelta(seconds=1)) update_license_ledger_for_automanaged_plan(realm, plan, self.next_year + timedelta(seconds=1))
ledger_entries = list(LicenseLedger.objects.values_list( ledger_entries = list(LicenseLedger.objects.values_list(
@ -1195,19 +1195,19 @@ class InvoiceTest(StripeTestCase):
with patch('corporate.lib.stripe.timezone_now', return_value=self.now): with patch('corporate.lib.stripe.timezone_now', return_value=self.now):
self.upgrade() self.upgrade()
# Increase # Increase
with patch('corporate.lib.stripe.get_seat_count', return_value=self.seat_count + 3): with patch('corporate.lib.stripe.get_latest_seat_count', return_value=self.seat_count + 3):
update_license_ledger_if_needed(get_realm('zulip'), self.now + timedelta(days=100)) update_license_ledger_if_needed(get_realm('zulip'), self.now + timedelta(days=100))
# Decrease # Decrease
with patch('corporate.lib.stripe.get_seat_count', return_value=self.seat_count): with patch('corporate.lib.stripe.get_latest_seat_count', return_value=self.seat_count):
update_license_ledger_if_needed(get_realm('zulip'), self.now + timedelta(days=200)) update_license_ledger_if_needed(get_realm('zulip'), self.now + timedelta(days=200))
# Increase, but not past high watermark # Increase, but not past high watermark
with patch('corporate.lib.stripe.get_seat_count', return_value=self.seat_count + 1): with patch('corporate.lib.stripe.get_latest_seat_count', return_value=self.seat_count + 1):
update_license_ledger_if_needed(get_realm('zulip'), self.now + timedelta(days=300)) update_license_ledger_if_needed(get_realm('zulip'), self.now + timedelta(days=300))
# Increase, but after renewal date, and below last year's high watermark # Increase, but after renewal date, and below last year's high watermark
with patch('corporate.lib.stripe.get_seat_count', return_value=self.seat_count + 2): with patch('corporate.lib.stripe.get_latest_seat_count', return_value=self.seat_count + 2):
update_license_ledger_if_needed(get_realm('zulip'), self.now + timedelta(days=400)) update_license_ledger_if_needed(get_realm('zulip'), self.now + timedelta(days=400))
# Increase, but after event_time # Increase, but after event_time
with patch('corporate.lib.stripe.get_seat_count', return_value=self.seat_count + 3): with patch('corporate.lib.stripe.get_latest_seat_count', return_value=self.seat_count + 3):
update_license_ledger_if_needed(get_realm('zulip'), self.now + timedelta(days=500)) update_license_ledger_if_needed(get_realm('zulip'), self.now + timedelta(days=500))
plan = CustomerPlan.objects.first() plan = CustomerPlan.objects.first()
invoice_plan(plan, self.now + timedelta(days=400)) invoice_plan(plan, self.now + timedelta(days=400))

View File

@ -16,7 +16,7 @@ from zerver.lib.response import json_error, json_success
from zerver.lib.validator import check_string, check_int from zerver.lib.validator import check_string, check_int
from zerver.models import UserProfile from zerver.models import UserProfile
from corporate.lib.stripe import STRIPE_PUBLISHABLE_KEY, \ from corporate.lib.stripe import STRIPE_PUBLISHABLE_KEY, \
stripe_get_customer, get_seat_count, \ stripe_get_customer, get_latest_seat_count, \
process_initial_upgrade, sign_string, \ process_initial_upgrade, sign_string, \
unsign_string, BillingError, do_change_plan_status, do_replace_payment_source, \ unsign_string, BillingError, do_change_plan_status, do_replace_payment_source, \
MIN_INVOICED_LICENSES, DEFAULT_INVOICE_DAYS_UNTIL_DUE, \ MIN_INVOICED_LICENSES, DEFAULT_INVOICE_DAYS_UNTIL_DUE, \
@ -124,7 +124,7 @@ def initial_upgrade(request: HttpRequest) -> HttpResponse:
if customer is not None and customer.default_discount is not None: if customer is not None and customer.default_discount is not None:
percent_off = customer.default_discount percent_off = customer.default_discount
seat_count = get_seat_count(user.realm) seat_count = get_latest_seat_count(user.realm)
signed_seat_count, salt = sign_string(str(seat_count)) signed_seat_count, salt = sign_string(str(seat_count))
context = { context = {
'publishable_key': STRIPE_PUBLISHABLE_KEY, 'publishable_key': STRIPE_PUBLISHABLE_KEY,
@ -177,7 +177,7 @@ def billing_home(request: HttpRequest) -> HttpResponse:
last_ledger_entry = make_end_of_cycle_updates_if_needed(plan, now) last_ledger_entry = make_end_of_cycle_updates_if_needed(plan, now)
if last_ledger_entry is not None: if last_ledger_entry is not None:
licenses = last_ledger_entry.licenses licenses = last_ledger_entry.licenses
licenses_used = get_seat_count(user.realm) licenses_used = get_latest_seat_count(user.realm)
# Should do this in javascript, using the user's timezone # Should do this in javascript, using the user's timezone
renewal_date = '{dt:%B} {dt.day}, {dt.year}'.format(dt=start_of_next_billing_cycle(plan, now)) renewal_date = '{dt:%B} {dt.day}, {dt.year}'.format(dt=start_of_next_billing_cycle(plan, now))
renewal_cents = renewal_amount(plan, now) renewal_cents = renewal_amount(plan, now)