Update endpoints to handle Enterprise plans

kobo/apps/organizations/serializers.py

@@ -1,13 +1,35 @@
 from rest_framework import serializers
 
-from kobo.apps.organizations.models import Organization, create_organization
+from kobo.apps.organizations.models import (
+    create_organization,
+    Organization,
+    OrganizationOwner,
+    OrganizationUser,
+)
+
+
+class OrganizationUserSerializer(serializers.ModelSerializer):
+
+    class Meta:
+        model = OrganizationUser
+        fields = ['user', 'organization']
+
+
+class OrganizationOwnerSerializer(serializers.ModelSerializer):
+    organization_user = OrganizationUserSerializer()
+
+    class Meta:
+        model = OrganizationOwner
+        fields = ['organization_user']
 
 
 class OrganizationSerializer(serializers.ModelSerializer):
+    owner = serializers.CharField(source='owner.organization_user.user.username', read_only=True)
+
     class Meta:
         model = Organization
-        fields = ['id', 'name', 'is_active', 'created', 'modified', 'slug']
-        read_only_fields = ["id", "slug"]
+        fields = ['id', 'name', 'is_active', 'created', 'modified', 'slug', 'owner']
+        read_only_fields = ['id', 'slug', 'owner']
 
     def create(self, validated_data):
         user = self.context['request'].user

kobo/apps/organizations/tests/test_organizations_api.py

@@ -44,7 +44,7 @@ def test_list(self):
         self._insert_data()
         organization2 = baker.make(Organization, id='org_abcd123')
         organization2.add_user(user=self.user, is_admin=True)
-        with self.assertNumQueries(FuzzyInt(2, 4)):
+        with self.assertNumQueries(FuzzyInt(8, 10)):
             res = self.client.get(self.url_list)
         self.assertContains(res, organization2.name)
 
@@ -63,7 +63,7 @@ def test_api_returns_org_data(self):
     def test_update(self):
         self._insert_data()
         data = {'name': 'edit'}
-        with self.assertNumQueries(FuzzyInt(4, 6)):
+        with self.assertNumQueries(FuzzyInt(8, 10)):
             res = self.client.patch(self.url_detail, data)
         self.assertContains(res, data['name'])
 

kobo/apps/stripe/serializers.py

@@ -100,7 +100,6 @@ class CheckoutLinkSerializer(PriceIdSerializer):
 
 
 class PriceSerializer(BasePriceSerializer):
-    product = BaseProductSerializer()
 
     class Meta(BasePriceSerializer.Meta):
         fields = (
@@ -114,18 +113,23 @@ class Meta(BasePriceSerializer.Meta):
             'metadata',
             'active',
             'product',
+            'transform_quantity',
         )
 
 
+class PriceWithProductSerializer(PriceSerializer):
+    product = BaseProductSerializer()
+
+
 class ProductSerializer(BaseProductSerializer):
-    prices = BasePriceSerializer(many=True)
+    prices = PriceSerializer(many=True)
 
     class Meta(BaseProductSerializer.Meta):
         fields = ('id', 'name', 'description', 'type', 'prices', 'metadata')
 
 
 class SubscriptionItemSerializer(serializers.ModelSerializer):
-    price = PriceSerializer()
+    price = PriceWithProductSerializer()
 
     class Meta:
         model = SubscriptionItem

kobo/apps/stripe/tests/test_organization_usage.py

@@ -0,0 +1,154 @@
+import timeit
+
+from dateutil.relativedelta import relativedelta
+from django.contrib.auth.models import User
+from django.core.cache import cache
+from django.urls import reverse
+from django.utils import timezone
+from djstripe.models import Customer, Price, Product, Subscription, SubscriptionItem
+from model_bakery import baker
+
+from kobo.apps.organizations.models import Organization, OrganizationUser
+from kobo.apps.trackers.submission_utils import create_mock_assets, add_mock_submissions
+from kpi.tests.api.v2.test_api_service_usage import ServiceUsageAPIBase
+
+
+class OrganizationUsageAPITestCase(ServiceUsageAPIBase):
+    """
+    Test organization service usage when Stripe is enabled.
+
+    Note: this class lives here (despite testing the Organizations endpoint) so that it will *only* run
+    when Stripe is installed.
+    """
+
+    user_count = 5
+    assets_per_user = 5
+    submissions_per_asset = 5
+    org_id = 'orgAKWMFskafsngf'
+
+    @classmethod
+    def setUpTestData(cls):
+        super().setUpTestData()
+        cls.now = timezone.now()
+
+        anotheruser = User.objects.get(username='anotheruser')
+        organization = baker.make(Organization, id=cls.org_id, name='test organization')
+        organization.add_user(cls.anotheruser, is_admin=True)
+        assets = create_mock_assets([cls.anotheruser], cls.assets_per_user)
+
+        cls.customer = baker.make(Customer, subscriber=organization, livemode=False)
+        organization.save()
+
+        users = baker.make(User, _quantity=cls.user_count - 1, _bulk_create=True)
+        baker.make(
+            OrganizationUser,
+            user=users.__iter__(),
+            organization=organization,
+            is_admin=False,
+            _quantity=cls.user_count - 1,
+            _bulk_create=True,
+        )
+        assets = assets + create_mock_assets(users, cls.assets_per_user)
+        add_mock_submissions(assets, cls.submissions_per_asset)
+
+    def setUp(self):
+        super().setUp()
+        url = reverse(self._get_endpoint('organizations-list'))
+        self.detail_url = f'{url}{self.org_id}/service_usage/'
+        self.expected_submissions_single = self.assets_per_user * self.submissions_per_asset
+        self.expected_submissions_multi = self.expected_submissions_single * self.user_count
+
+    def tearDown(self):
+        cache.clear()
+
+    def generate_subscription(self, metadata: dict):
+        """Create a subscription for a product with custom"""
+        product = baker.make(Product, active=True, metadata={
+            'product_type': 'plan',
+            **metadata,
+        })
+        price = baker.make(
+            Price,
+            active=True,
+            id='price_sfmOFe33rfsfd36685657',
+            product=product,
+        )
+
+        subscription_item = baker.make(SubscriptionItem, price=price, quantity=1, livemode=False)
+        baker.make(
+            Subscription,
+            customer=self.customer,
+            status='active',
+            items=[subscription_item],
+            livemode=False,
+            billing_cycle_anchor=self.now - relativedelta(weeks=2),
+            current_period_end=self.now + relativedelta(weeks=2),
+            current_period_start=self.now - relativedelta(weeks=2),
+        )
+
+    def test_usage_doesnt_include_org_users_without_subscription(self):
+        """
+        Test that the endpoint *only* returns usage for the logged-in user
+        if they don't have a subscription that includes Organizations.
+        """
+        response = self.client.get(self.detail_url)
+        # without a plan, the user should only see their usage
+        assert response.data['total_submission_count']['all_time'] == self.expected_submissions_single
+        assert response.data['total_submission_count']['current_month'] == self.expected_submissions_single
+        assert response.data['total_storage_bytes'] == (
+            self.expected_file_size() * self.expected_submissions_single
+        )
+
+    def test_usage_for_plans_with_org_access(self):
+        """
+        Test that the endpoint aggregates usage for each user in the organization
+        when viewing /service_usage/{organization_id}/
+        """
+
+        self.generate_subscription(
+            {
+                'plan_type': 'enterprise',
+                'organizations': True,
+            }
+        )
+
+        # the user should see usage for everyone in their org
+        response = self.client.get(self.detail_url)
+        assert response.data['total_submission_count']['current_month'] == self.expected_submissions_multi
+        assert response.data['total_submission_count']['all_time'] == self.expected_submissions_multi
+        assert response.data['total_storage_bytes'] == (
+            self.expected_file_size() * self.expected_submissions_multi
+        )
+
+    def test_doesnt_include_org_users_with_invalid_plan(self):
+        """
+        Test that the endpoint *doesn't* aggregates usage for the organization
+        when subscribed to a product that doesn't include org access
+        """
+
+        self.generate_subscription({})
+
+        response = self.client.get(self.detail_url)
+        # without the proper subscription, the user should only see their usage
+        assert response.data['total_submission_count']['current_month'] == self.expected_submissions_single
+        assert response.data['total_submission_count']['all_time'] == self.expected_submissions_single
+        assert response.data['total_storage_bytes'] == (
+            self.expected_file_size() * self.expected_submissions_single
+        )
+
+    def test_endpoint_speed_(self):
+        # get the average request time for 10 hits to the endpoint
+        single_user_time = timeit.timeit(lambda: self.client.get(self.detail_url), number=10)
+
+        self.generate_subscription(
+            {
+                'plan_type': 'enterprise',
+                'organizations': True,
+            }
+        )
+
+        # get the average request time for 10 hits to the endpoint
+        multi_user_time = timeit.timeit(lambda: self.client.get(self.detail_url), number=10)
+        assert single_user_time < 1.5
+        assert multi_user_time < 2
+        assert multi_user_time < single_user_time * 2

kobo/apps/stripe/views.py

@@ -324,14 +324,17 @@ def generate_portal_link(user, organization_id, price):
             if not len(all_configs):
                 return Response({'error': "Missing Stripe billing configuration."}, status=status.HTTP_502_BAD_GATEWAY)
 
-            is_price_for_addon = price.product.metadata.get('product_type', '') == 'addon'
-
-            if is_price_for_addon:
-                """
-                Recurring add-ons aren't included in the default billing configuration.
-                This lets us hide them as an 'upgrade' option for paid plan users.
-                Here, we try getting the portal configuration that lets us switch to the provided price.
-                """
+            """
+            Recurring add-ons and the Enterprise plan aren't included in the default billing configuration.
+            This lets us hide them as an 'upgrade' option for paid plan users.
+            """
+            metadata = price.product.metadata
+            needs_custom_config = (
+                metadata.get('product_type') == 'addon'
+                or metadata.get('plan_type') == 'enterprise'
+            )
+            if needs_custom_config:
+                # Try getting the portal configuration that lets us switch to the provided price
                 current_config = next(
                     (config for config in all_configs if (
                             config['active'] and
@@ -350,7 +353,7 @@ def generate_portal_link(user, organization_id, price):
                     )), None
                 )
 
-                if is_price_for_addon:
+                if needs_custom_config:
                     """
                     we couldn't find a custom configuration, let's try making a new one
                     add the price we're switching into to the list of prices that allow subscription updates
@@ -474,7 +477,10 @@ class ProductViewSet(viewsets.GenericViewSet, mixins.ListModelMixin):
     >                           },
     >                           "unit_amount": int (cents),
     >                           "human_readable_price": string,
-    >                           "metadata": {}
+    >                           "metadata": {},
+    >                           "active": bool,
+    >                           "product": string,
+    >                           "transform_quantity": null | {'round': 'up'|'down', 'divide_by': int}
     >                       },
     >                       ...
     >                   ],