Add unit tests for new aggregation annotation approach

InvenTree/company/fixtures/company.yaml

@@ -14,4 +14,11 @@
   pk: 3
   fields:
     name: Zerg Corp
-    description: We eat the competition
+    description: We eat the competition
+- model: company.company
+  pk: 4
+  fields:
+    name: A customer
+    description: A company that we sell things to!
+    is_customer: True
+  

InvenTree/part/serializers.py

@@ -15,7 +15,7 @@ from .models import PartTestTemplate
 
 from decimal import Decimal
 
-from sql_util.utils import SubquerySum
+from sql_util.utils import SubquerySum, SubqueryCount
 
 from django.db.models import Q
 from django.db.models.functions import Coalesce
@@ -208,6 +208,11 @@ class PartSerializer(InvenTreeModelSerializer):
             ),
         )
 
+        # Annotate with the total number of stock items
+        queryset = queryset.annotate(
+            stock_item_count=SubqueryCount('stock_items')
+        )
+
         # Filter to limit builds to "active"
         build_filter = Q(
             status__in=BuildStatus.ACTIVE_CODES
@@ -253,6 +258,7 @@ class PartSerializer(InvenTreeModelSerializer):
     in_stock = serializers.FloatField(read_only=True)
     ordering = serializers.FloatField(read_only=True)
     building = serializers.FloatField(read_only=True)
+    stock_item_count = serializers.IntegerField(read_only=True)
 
     image = serializers.CharField(source='get_image_url', read_only=True)
     thumbnail = serializers.CharField(source='get_thumbnail_url', read_only=True)
@@ -295,6 +301,7 @@ class PartSerializer(InvenTreeModelSerializer):
             'revision',
             'salable',
             'starred',
+            'stock_item_count',
             'thumbnail',
             'trackable',
             'units',

InvenTree/part/test_api.py

@@ -1,8 +1,15 @@
 from rest_framework.test import APITestCase
 from rest_framework import status
+
 from django.urls import reverse
 from django.contrib.auth import get_user_model
 
+from part.models import Part
+from stock.models import StockItem
+from company.models import Company
+
+from InvenTree.status_codes import StockStatus
+
 
 class PartAPITest(APITestCase):
     """
@@ -213,3 +220,77 @@ class PartAPITest(APITestCase):
         )
 
         self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
+
+
+class PartAPIAggregationTest(APITestCase):
+    """
+    Tests to ensure that the various aggregation annotations are working correctly...
+    """
+
+    fixtures = [
+        'category',
+        'company',
+        'part',
+        'location',
+        'bom',
+        'test_templates',
+    ]
+
+    def setUp(self):
+        # Create a user for auth
+        User = get_user_model()
+        User.objects.create_user('testuser', 'test@testing.com', 'password')
+
+        self.client.login(username='testuser', password='password')
+
+        # Add a new part
+        self.part = Part.objects.create(
+            name='Banana',
+        )
+
+        # Create some stock items associated with the part
+
+        # First create 600 units which are OK
+        StockItem.objects.create(part=self.part, quantity=100)
+        StockItem.objects.create(part=self.part, quantity=200)
+        StockItem.objects.create(part=self.part, quantity=300)
+
+        # Now create another 400 units which are LOST
+        StockItem.objects.create(part=self.part, quantity=400, status=StockStatus.LOST)
+
+    def get_part_data(self):
+        url = reverse('api-part-list')
+
+        response = self.client.get(url, format='json')
+
+        self.assertEqual(response.status_code, status.HTTP_200_OK)
+
+        for part in response.data:
+            if part['pk'] == self.part.pk:
+                return part
+
+        # We should never get here!
+        self.assertTrue(False)
+
+    def test_stock_quantity(self):
+        """
+        Simple test for the stock quantity
+        """
+
+        data = self.get_part_data()
+
+        self.assertEqual(data['in_stock'], 600)
+        self.assertEqual(data['stock_item_count'], 4)
+    
+        # Add some more stock items!!
+        for i in range(100):
+            StockItem.objects.create(part=self.part, quantity=5)
+
+        # Add another stock item which is assigned to a customer (and shouldn't count)
+        customer = Company.objects.get(pk=4)
+        StockItem.objects.create(part=self.part, quantity=9999, customer=customer)
+
+        data = self.get_part_data()
+
+        self.assertEqual(data['in_stock'], 1100)
+        self.assertEqual(data['stock_item_count'], 105)