From d27d0c333347e3d55b0cbdbf70a9f63b0263bde4 Mon Sep 17 00:00:00 2001
From: Valentin Rigal <rigal@teklia.com>
Date: Wed, 27 Sep 2023 13:32:36 +0200
Subject: [PATCH] Update Serializer usages

---
 tests/push_health/test_compare.py | 1 +
 treeherder/push_health/usage.py   | 5 ++++-
 treeherder/webapp/api/push.py     | 8 ++++++--
 3 files changed, 11 insertions(+), 3 deletions(-)

diff --git a/tests/push_health/test_compare.py b/tests/push_health/test_compare.py
index f17f9ef026c..414b1d25755 100644
--- 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
index d6cffabcb8d..9fd056a94b8 100644
--- 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
index 3667f9e4055..62a1e652bfb 100644
--- 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