From 547e976cbac2e5f1bf49a6165cf8ba4c6a435e99 Mon Sep 17 00:00:00 2001
From: Shivank Kacker <kacker.shivank@gmail.com>
Date: Wed, 9 Oct 2024 13:18:59 +0530
Subject: [PATCH] Added /hubs route to list hubs

---
 care/facility/api/serializers/facility.py | 11 -----------
 care/facility/api/viewsets/facility.py    | 16 ++++++++++++++++
 care/facility/tests/test_facility_api.py  | 19 +++++++++++++++++++
 config/api_router.py                      |  2 ++
 4 files changed, 37 insertions(+), 11 deletions(-)

diff --git a/care/facility/api/serializers/facility.py b/care/facility/api/serializers/facility.py
index 6cceceb9db..bd42b9bba4 100644
--- a/care/facility/api/serializers/facility.py
+++ b/care/facility/api/serializers/facility.py
@@ -105,18 +105,9 @@ class FacilitySerializer(FacilityBasicInfoSerializer):
 
     facility_flags = serializers.SerializerMethodField()
 
-    hubs = serializers.SerializerMethodField()
-    spokes = serializers.SerializerMethodField()
-
     def get_facility_flags(self, facility):
         return facility.get_facility_flags()
 
-    def get_hubs(self, facility):
-        return FacilitySpokeSerializer(facility.hubs, many=True).data
-
-    def get_spokes(self, facility):
-        return FacilitySpokeSerializer(facility.spokes, many=True).data
-
     class Meta:
         model = Facility
         fields = [
@@ -153,8 +144,6 @@ class Meta:
             "patient_count",
             "bed_count",
             "facility_flags",
-            "hubs",
-            "spokes",
         ]
         read_only_fields = ("modified_date", "created_date")
 
diff --git a/care/facility/api/viewsets/facility.py b/care/facility/api/viewsets/facility.py
index dd8221c0cb..f6598a539c 100644
--- a/care/facility/api/viewsets/facility.py
+++ b/care/facility/api/viewsets/facility.py
@@ -204,3 +204,19 @@ def get_serializer_context(self):
         context = super().get_serializer_context()
         context["facility"] = facility
         return context
+
+
+class FacilityHubsViewSet(mixins.ListModelMixin, viewsets.GenericViewSet):
+    queryset = FacilityHubSpoke.objects.all().select_related("spoke", "hub")
+    serializer_class = FacilitySpokeSerializer
+    permission_classes = (IsAuthenticated, DRYPermissions)
+    lookup_field = "external_id"
+
+    def get_queryset(self):
+        return self.queryset.filter(spoke=self.get_facility())
+
+    def get_facility(self):
+        facilities = get_facility_queryset(self.request.user)
+        return get_object_or_404(
+            facilities.filter(external_id=self.kwargs["facility_external_id"])
+        )
diff --git a/care/facility/tests/test_facility_api.py b/care/facility/tests/test_facility_api.py
index 800f45fb8e..b86cb0db86 100644
--- a/care/facility/tests/test_facility_api.py
+++ b/care/facility/tests/test_facility_api.py
@@ -229,6 +229,25 @@ def test_spoke_is_not_ancestor(self):
         )
         self.assertIs(response.status_code, status.HTTP_400_BAD_REQUEST)
 
+    def test_hubs_list(self):
+        facility_a = self.create_facility(
+            self.super_user, self.district, self.local_body
+        )
+        facility_b = self.create_facility(
+            self.super_user, self.district, self.local_body
+        )
+
+        FacilityHubSpoke.objects.create(hub=facility_a, spoke=facility_b)
+
+        self.client.force_authenticate(user=self.super_user)
+        response = self.client.get(f"/api/v1/facility/{facility_b.external_id}/hubs/")
+        self.assertIs(response.status_code, status.HTTP_200_OK)
+        data = response.json()
+        self.assertEqual(data["count"], 1)
+        self.assertEqual(
+            data["results"][0]["hub_object"]["id"], str(facility_a.external_id)
+        )
+
 
 class FacilityCoverImageTests(TestUtils, APITestCase):
     @classmethod
diff --git a/config/api_router.py b/config/api_router.py
index 917b187395..a2e3de78c4 100644
--- a/config/api_router.py
+++ b/config/api_router.py
@@ -37,6 +37,7 @@
 )
 from care.facility.api.viewsets.facility import (
     AllFacilityViewSet,
+    FacilityHubsViewSet,
     FacilitySpokesViewSet,
     FacilityViewSet,
 )
@@ -218,6 +219,7 @@
 facility_nested_router.register(
     r"spokes", FacilitySpokesViewSet, basename="facility-spokes"
 )
+facility_nested_router.register(r"hubs", FacilityHubsViewSet, basename="facility-hubs")
 
 router.register("asset", AssetViewSet, basename="asset")
 asset_nested_router = NestedSimpleRouter(router, r"asset", lookup="asset")