[Do not merge] Testing codecov AI flows

graphs/urls.py

@@ -1,6 +1,6 @@
 from django.urls import re_path
 
-from .views import BadgeHandler, GraphHandler
+from .views import BadgeHandler, BundleBadgeHandler, GraphHandler
 
 urlpatterns = [
     re_path(
@@ -13,6 +13,16 @@
         BadgeHandler.as_view(),
         name="default-badge",
     ),
+    re_path(
+        "branch/(?P<branch>.+)/(graph|graphs)/bundle/(?P<bundle>.+)/badge.(?P<ext>[^/]+)",
+        BundleBadgeHandler.as_view(),
+        name="branch-bundle-badge",
+    ),
+    re_path(
+        "(graph|graphs)/bundle/(?P<bundle>.+)/badge.(?P<ext>[^/]+)",
+        BundleBadgeHandler.as_view(),
+        name="default-bundle-badge",
+    ),
     re_path(
         "pull/(?P<pullid>[^/]+)/(graph|graphs)/(?P<graph>tree|icicle|sunburst|commits).(?P<ext>[^/]+)",
         GraphHandler.as_view(),

graphs/views.py

@@ -14,9 +14,15 @@
 from api.shared.mixins import RepoPropertyMixin
 from core.models import Branch, Pull
 from graphs.settings import settings
+from services.bundle_analysis import load_report
 from services.components import commit_components
 
-from .helpers.badge import format_coverage_precision, get_badge
+from .helpers.badge import (
+    format_bundle_bytes,
+    format_coverage_precision,
+    get_badge,
+    get_bundle_badge,
+)
 from .helpers.graphs import icicle, sunburst, tree
 from .mixins import GraphBadgeAPIMixin
 
@@ -193,6 +199,86 @@
         return filtered_report.totals.coverage
 
 
+class BundleBadgeHandler(APIView, RepoPropertyMixin, GraphBadgeAPIMixin):
+    content_negotiation_class = IgnoreClientContentNegotiation
+
+    permission_classes = [AllowAny]
+
+    extensions = ["svg", "txt"]
+    precisions = ["0", "1", "2"]
+    filename = "bundle-badge"
+
+    def get_object(self, request, *args, **kwargs):
+        # Validate precision query param
+        precision = self.request.query_params.get("precision", "2")
+        if precision not in self.precisions:
+            raise NotFound("Bundle size precision should be one of [ 0 || 1 || 2 ]")
+        precision = int(precision)
+
+        bundle_size_bytes = self.get_bundle_size()
+
+        if self.kwargs.get("ext") == "txt":
+            return (
+                "unknown"
+                if bundle_size_bytes is None
+                else format_bundle_bytes(bundle_size_bytes, precision)
+            )
+
+        return get_bundle_badge(bundle_size_bytes, precision)
+
+    def get_bundle_size(self) -> int | None:
+        try:
+            repo = self.repo
+        except Http404:
+            log.warning("Repo not found", extra=dict(repo=self.kwargs.get("repo_name")))
+            return None
+
+        if repo.private and repo.image_token != self.request.query_params.get("token"):
+            log.warning(
+                "Token provided does not match repo's image token",
+                extra=dict(repo=repo),
+            )
+            return None
+
+        branch_name = self.kwargs.get("branch") or repo.branch
+        branch = Branch.objects.filter(
+            name=branch_name, repository_id=repo.repoid
+        ).first()
+
+        if branch is None:
+            log.warning(
+                "Branch not found", extra=dict(branch_name=branch_name, repo=repo)
+            )
+            return None
+
+        try:
+            commit: Commit = repo.commits.filter(commitid=branch.head).first()
+        except ObjectDoesNotExist:
+            log.warning("Commit not found", extra=dict(commit=branch.head))
+            return None
+
+        commit_bundles = load_report(commit)
+
+        if commit_bundles is None:
+            log.warning(
+                "Bundle analysis report not found for commit",
+                extra=dict(commit=branch.head),
+            )
+            return None
+
+        bundle_name = str(self.kwargs.get("bundle"))
+        bundle = commit_bundles.bundle_report(bundle_name)
+
+        if bundle is None:
+            log.warning(
+                "Bundle with provided name not found for commit",
+                extra=dict(commit=branch.head),
+            )
+            return None
+
+        return bundle.total_size()
+
+
 class GraphHandler(APIView, RepoPropertyMixin, GraphBadgeAPIMixin):
     permission_classes = [AllowAny]