Update Serializer usages

vrigal · Sep 28, 2023 · d27d0c3 · d27d0c3
1 parent e740331
commit d27d0c3
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 3 deletions.
diff --git a/tests/push_health/test_compare.py b/tests/push_health/test_compare.py
@@ -103,6 +103,7 @@ def test_get_commit_history(test_push, test_repository, mock_rev, mock_json_push
         author='foo@bar.baz',
         time=datetime.datetime.now(),
     )
+    test_push.revision_count = test_push.commits.count()
 
     history = get_commit_history(test_repository, test_revision, test_push)
     print('\n<><><>history')

diff --git a/treeherder/push_health/usage.py b/treeherder/push_health/usage.py
@@ -1,5 +1,6 @@
 import logging
 
+from django.db.models import Count
 from treeherder.config import settings
 from treeherder.model.models import Push, Job
 from treeherder.push_health.classification import NEED_INVESTIGATION
@@ -55,7 +56,9 @@ def get_usage():
 
     results = [
         {
-            'push': PushSerializer(pushes.get(revision=facet['name'])).data,
+            'push': PushSerializer(
+                pushes.annotate(revision_count=Count('commits')).get(revision=facet['name'])
+            ).data,
             'peak': get_peak(facet),
             'latest': get_latest(facet),
             'retriggers': jobs_retriggered(pushes.get(revision=facet['name'])),

diff --git a/treeherder/webapp/api/push.py b/treeherder/webapp/api/push.py
@@ -197,7 +197,9 @@ def retrieve(self, request, project, pk=None):
         GET method implementation for detail view of ``push``
         """
         try:
-            push = Push.objects.get(repository__name=project, id=pk)
+            push = Push.objects.annotate(revision_count=Count('commits')).get(
+                repository__name=project, id=pk
+            )
             serializer = PushSerializer(push)
             return Response(serializer.data)
         except Push.DoesNotExist:
@@ -340,7 +342,9 @@ def health(self, request, project):
 
         try:
             repository = Repository.objects.get(name=project)
-            push = Push.objects.get(revision=revision, repository=repository)
+            push = Push.objects.annotate(revision_count=Count('commits')).get(
+                revision=revision, repository=repository
+            )
         except Push.DoesNotExist:
             return Response(
                 "No push with revision: {0}".format(revision), status=HTTP_404_NOT_FOUND