Merge pull request #5052 from kobotoolbox/TASK-934-update-canceled-sub-handling-backend

[TASK-934] Anchor default plan cycle to canceled subscription end date

Author: jamesrkiger

kobo/apps/organizations/models.py

@@ -51,6 +51,26 @@ def active_subscription_billing_details(self):
                 ).first()
 
         return None
+    
+    @cache_for_request
+    def canceled_subscription_billing_cycle_anchor(self):
+        """
+        Returns cancelation date of most recently canceled subscription
+        """
+        # Only check for subscriptions if Stripe is enabled
+        if settings.STRIPE_ENABLED:
+            qs = Organization.objects.prefetch_related('djstripe_customers').filter(
+                    djstripe_customers__subscriptions__status='canceled',
+                    djstripe_customers__subscriber=self.id,
+                ).order_by(
+                    '-djstripe_customers__subscriptions__ended_at'
+                ).values(
+                    anchor=F('djstripe_customers__subscriptions__ended_at'),
+                ).first()
+            if qs:
+                return qs['anchor']
+            
+        return None
 
 
 class OrganizationUser(AbstractOrganizationUser):

kobo/apps/organizations/utils.py

@@ -7,43 +7,79 @@
 from kobo.apps.organizations.models import Organization
 
 
-def organization_month_start(organization: Union[Organization, None]):
-    now = timezone.now()
-    first_of_this_month = now.date().replace(day=1)
-    # If no organization/subscription, just use the first day of current month
+def get_monthly_billing_dates(organization: Union[Organization, None]):
+    """Returns start and end dates of an organization's monthly billing cycle"""
+
+    now = timezone.now().replace(tzinfo=pytz.UTC)
+    first_of_this_month = now.replace(day=1)
+    # Using `days=31` gets the last day of the month
+    last_of_this_month = first_of_this_month + relativedelta(days=31)
+
+    # If no organization, just use the calendar month
     if not organization:
-        return first_of_this_month
+        return first_of_this_month, last_of_this_month
+    
+    # If no active subscription, check for canceled subscription 
     if not (billing_details := organization.active_subscription_billing_details()):
-        return first_of_this_month
+        if not (
+            canceled_subscription_anchor
+            := organization.canceled_subscription_billing_cycle_anchor()
+        ):
+            return first_of_this_month, last_of_this_month
+        
+        period_end = canceled_subscription_anchor.replace(tzinfo=pytz.UTC)
+        while period_end < now:
+            period_end += relativedelta(months=1)
+        period_start = period_end - relativedelta(months=1)
+        return period_start, period_end
+    
     if not billing_details.get('billing_cycle_anchor'):
-        return first_of_this_month
+        return first_of_this_month, last_of_this_month
 
-    # Subscription is billed monthly, use the current billing period start date
+    # Subscription is billed monthly, use the current billing period dates
     if billing_details.get('recurring_interval') == 'month':
-        return billing_details.get('current_period_start').replace(tzinfo=pytz.UTC)
+        period_start = billing_details.get('current_period_start').replace(
+            tzinfo=pytz.UTC
+        )
+        period_end = billing_details.get('current_period_end').replace(
+            tzinfo=pytz.UTC
+        )
+        return period_start, period_end
 
     # Subscription is billed yearly - count backwards from the end of the current billing year
-    month_start = billing_details.get('current_period_end').replace(tzinfo=pytz.UTC)
-    while month_start > now:
-        month_start -= relativedelta(months=1)
-    return month_start
+    period_start = billing_details.get('current_period_end').replace(tzinfo=pytz.UTC)
+    while period_start > now:
+        period_start -= relativedelta(months=1)
+    period_end = period_start + relativedelta(months=1)
+    return period_start, period_end
 
 
-def organization_year_start(organization: Union[Organization, None]):
-    now = timezone.now()
+def get_yearly_billing_dates(organization: Union[Organization, None]):
+    """Returns start and end dates of an organization's annual billing cycle"""
+    now = timezone.now().replace(tzinfo=pytz.UTC)
     first_of_this_year = now.date().replace(month=1, day=1)
+    last_of_this_year = now.date().replace(month=12, day=31)
+
     if not organization:
-        return first_of_this_year
+        return first_of_this_year, last_of_this_year
     if not (billing_details := organization.active_subscription_billing_details()):
-        return first_of_this_year
+        return first_of_this_year, last_of_this_year
     if not (anchor_date := billing_details.get('billing_cycle_anchor')):
-        return first_of_this_year
+        return first_of_this_year, last_of_this_year
 
-    # Subscription is billed yearly, use the provided anchor date as start date
+    # Subscription is billed yearly, use the dates from the subscription
     if billing_details.get('subscription_interval') == 'year':
-        return billing_details.get('current_period_start').replace(tzinfo=pytz.UTC)
+        period_start = billing_details.get('current_period_start').replace(
+            tzinfo=pytz.UTC
+        )
+        period_end = billing_details.get('current_period_end').replace(
+            tzinfo=pytz.UTC
+        )
+        return period_start, period_end
 
     # Subscription is monthly, calculate this year's start based on anchor date
-    while anchor_date + relativedelta(years=1) < now:
+    period_start = anchor_date.replace(tzinfo=pytz.UTC) + relativedelta(years=1)
+    while period_start < now:
         anchor_date += relativedelta(years=1)
-    return anchor_date.replace(tzinfo=pytz.UTC)
+    period_end = period_start + relativedelta(years=1)
+    return period_start, period_end

kobo/apps/project_ownership/tests/api/v2/test_api.py

@@ -394,9 +394,6 @@ def test_account_usage_transferred_to_new_user(self):
                 'current_year': 1,
                 'current_month': 1,
             },
-            'current_month_start': today.replace(day=1).strftime('%Y-%m-%d'),
-            'current_year_start': today.replace(month=1, day=1).strftime('%Y-%m-%d'),
-            'billing_period_end': None,
         }
 
         expected_empty_data = {
@@ -414,9 +411,6 @@ def test_account_usage_transferred_to_new_user(self):
                 'current_year': 0,
                 'current_month': 0,
             },
-            'current_month_start': today.replace(day=1).strftime('%Y-%m-%d'),
-            'current_year_start': today.replace(month=1, day=1).strftime('%Y-%m-%d'),
-            'billing_period_end': None,
         }
 
         service_usage_url = reverse(
@@ -425,12 +419,16 @@ def test_account_usage_transferred_to_new_user(self):
         # someuser has some usage metrics
         self.client.login(username='someuser', password='someuser')
         response = self.client.get(service_usage_url)
-        assert response.data == expected_data
+        assert response.data['total_nlp_usage'] == expected_data['total_nlp_usage']
+        assert response.data['total_storage_bytes'] == expected_data['total_storage_bytes']
+        assert response.data['total_submission_count'] == expected_data['total_submission_count']
 
         # anotheruser's usage should be 0
         self.client.login(username='anotheruser', password='anotheruser')
         response = self.client.get(service_usage_url)
-        assert response.data == expected_empty_data
+        assert response.data['total_nlp_usage'] == expected_empty_data['total_nlp_usage']
+        assert response.data['total_storage_bytes'] == expected_empty_data['total_storage_bytes']
+        assert response.data['total_submission_count'] == expected_empty_data['total_submission_count']
 
         # Transfer project from someuser to anotheruser
         self.client.login(username='someuser', password='someuser')
@@ -452,12 +450,16 @@ def test_account_usage_transferred_to_new_user(self):
 
         # someuser should have no usage reported anymore
         response = self.client.get(service_usage_url)
-        assert response.data == expected_empty_data
+        assert response.data['total_nlp_usage'] == expected_empty_data['total_nlp_usage']
+        assert response.data['total_storage_bytes'] == expected_empty_data['total_storage_bytes']
+        assert response.data['total_submission_count'] == expected_empty_data['total_submission_count']
 
         # anotheruser should now have usage reported
         self.client.login(username='anotheruser', password='anotheruser')
         response = self.client.get(service_usage_url)
-        assert response.data == expected_data
+        assert response.data['total_nlp_usage'] == expected_data['total_nlp_usage']
+        assert response.data['total_storage_bytes'] == expected_data['total_storage_bytes']
+        assert response.data['total_submission_count'] == expected_data['total_submission_count']
 
     @patch(
         'kobo.apps.project_ownership.models.transfer.reset_kc_permissions',

kobo/apps/stripe/tests/test_organization_usage.py

@@ -1,6 +1,7 @@
 import timeit
 
 import pytest
+from dateutil.relativedelta import relativedelta
 from django.core.cache import cache
 from django.test import override_settings
 from django.urls import reverse
@@ -11,14 +12,14 @@
 from kobo.apps.kobo_auth.shortcuts import User
 from kobo.apps.organizations.models import Organization, OrganizationUser
 from kobo.apps.stripe.tests.utils import generate_enterprise_subscription, generate_plan_subscription
-from kobo.apps.trackers.submission_utils import create_mock_assets, add_mock_submissions
+from kobo.apps.trackers.tests.submission_utils import create_mock_assets, add_mock_submissions
 from kpi.tests.api.v2.test_api_service_usage import ServiceUsageAPIBase
 from kpi.tests.api.v2.test_api_asset_usage import AssetUsageAPITestCase
 from rest_framework import status
 
 
 
-class OrganizationServiceUsageAPITestCase(ServiceUsageAPIBase):
+class OrganizationServiceUsageAPIMultiUserTestCase(ServiceUsageAPIBase):
     """
     Test organization service usage when Stripe is enabled.
 
@@ -146,6 +147,83 @@ def test_endpoint_is_cached(self):
             self.expected_file_size() * self.expected_submissions_multi
         )
 
+class OrganizationServiceUsageAPITestCase(ServiceUsageAPIBase):
+    org_id = 'orgAKWMFskafsngf'
+
+    @classmethod
+    def setUpTestData(cls):
+        super().setUpTestData()
+        cls.now = timezone.now()
+
+        cls.organization = baker.make(
+            Organization, id=cls.org_id, name='test organization'
+        )
+        cls.organization.add_user(cls.anotheruser, is_admin=True)
+        cls.asset = create_mock_assets([cls.anotheruser])[0]
+
+
+    def setUp(self):
+        super().setUp()
+        url = reverse(self._get_endpoint('organizations-list'))
+        self.detail_url = f'{url}{self.org_id}/service_usage/'
+
+    def tearDown(self):
+        cache.clear()
+
+    def test_plan_canceled_this_month(self):
+        """
+        When a user cancels their subscription, they revert to the default community plan
+        with a billing cycle anchored to the end date of their canceled subscription
+        """
+
+        subscription = generate_plan_subscription(self.organization, age_days=30)
+
+        num_submissions = 5
+        add_mock_submissions([self.asset], num_submissions, 15)
+
+        canceled_at = timezone.now()
+        subscription.status = 'canceled'
+        subscription.ended_at = canceled_at
+        subscription.save()
+
+        current_billing_period_end = canceled_at + relativedelta(months=1)
+
+        response = self.client.get(self.detail_url)
+        
+        assert (
+            response.data['total_submission_count']['current_month']
+            == 0
+        )
+        assert response.data['current_month_start'] == canceled_at.isoformat()
+        assert response.data['current_month_end'] == current_billing_period_end.isoformat()
+
+    def test_plan_canceled_last_month(self):
+        subscription = generate_plan_subscription(self.organization, age_days=60)
+
+        num_submissions = 5
+        add_mock_submissions([self.asset], num_submissions, 15)
+
+        canceled_at = timezone.now() - relativedelta(days=45)
+        subscription.status = 'canceled'
+        subscription.ended_at = canceled_at
+        subscription.save()
+
+        current_billing_period_start = canceled_at + relativedelta(months=1)
+        current_billing_period_end = current_billing_period_start + relativedelta(
+            months=1
+        )
+        response = self.client.get(self.detail_url)
+        
+        assert (
+            response.data['total_submission_count']['current_month']
+            == 5
+        )
+        assert response.data['current_month_start'] == current_billing_period_start.isoformat()
+        assert (
+            response.data['current_month_end']
+            == current_billing_period_end.isoformat()
+        )
+
 
 class OrganizationAssetUsageAPITestCase(AssetUsageAPITestCase):
     """

kobo/apps/stripe/tests/utils.py

@@ -7,7 +7,7 @@
 
 
 def generate_plan_subscription(
-    organization: Organization, metadata: dict = None, customer: Customer = None
+    organization: Organization, metadata: dict = None, customer: Customer = None, age_days=0
 ) -> Subscription:
     """Create a subscription for a product with custom metadata"""
     now = timezone.now()

kobo/apps/trackers/submission_utils.py renamed to kobo/apps/trackers/tests/submission_utils.py

@@ -1,6 +1,7 @@
 import os
 import uuid
 
+from dateutil.relativedelta import relativedelta
 from django.conf import settings
 from django.utils import timezone
 from model_bakery import baker
@@ -61,12 +62,15 @@ def expected_file_size(submissions: int = 1):
 
 
 def update_xform_counters(
-    asset: Asset, xform: KobocatXForm = None, submissions: int = 1
+    asset: Asset,
+    xform: KobocatXForm = None,
+    submissions: int = 1,
+    age_days: int = 0,
 ):
     """
     Create/update the daily submission counter and the shadow xform we use to query it
     """
-    today = timezone.now()
+    today = timezone.now() - relativedelta(days=age_days)
     if xform:
         xform.attachment_storage_bytes += (
             expected_file_size(submissions)
@@ -124,7 +128,7 @@ def update_xform_counters(
         counter.save()
 
 
-def add_mock_submissions(assets: list, submissions_per_asset: int = 1):
+def add_mock_submissions(assets: list, submissions_per_asset: int = 1, age_days: int = 0):
     """
     Add one (default) or more submissions to an asset
     """
@@ -158,6 +162,6 @@ def add_mock_submissions(assets: list, submissions_per_asset: int = 1):
 
         asset.deployment.mock_submissions(asset_submissions, flush_db=False)
         all_submissions = all_submissions + asset_submissions
-        update_xform_counters(asset, submissions=submissions_per_asset)
+        update_xform_counters(asset, submissions=submissions_per_asset, age_days=age_days)
 
     return all_submissions