make asset, bed, assetbed get_queryset reusable based on user

ohcnetwork · Oct 15, 2024 · 3aa52e8 · 3aa52e8
1 parent 55b9eb8
commit 3aa52e8
Show file tree

Hide file tree

Showing 4 changed files with 65 additions and 96 deletions.
diff --git a/care/facility/api/viewsets/asset.py b/care/facility/api/viewsets/asset.py
@@ -62,6 +62,7 @@
 from care.utils.assetintegration.asset_classes import AssetClasses
 from care.utils.cache.cache_allowed_facilities import get_accessible_facilities
 from care.utils.filters.choicefilter import CareChoiceFilter, inverse_choices
+from care.utils.queryset.asset_bed import get_asset_queryset
 from care.utils.queryset.asset_location import get_asset_location_queryset
 from care.utils.queryset.facility import get_facility_queryset
 from config.authentication import MiddlewareAuthentication
@@ -290,21 +291,7 @@ class AssetViewSet(
     filterset_class = AssetFilter
 
     def get_queryset(self):
-        user = self.request.user
-        queryset = self.queryset
-        if user.is_superuser:
-            pass
-        elif user.user_type >= User.TYPE_VALUE_MAP["StateLabAdmin"]:
-            queryset = queryset.filter(current_location__facility__state=user.state)
-        elif user.user_type >= User.TYPE_VALUE_MAP["DistrictLabAdmin"]:
-            queryset = queryset.filter(
-                current_location__facility__district=user.district
-            )
-        else:
-            allowed_facilities = get_accessible_facilities(user)
-            queryset = queryset.filter(
-                current_location__facility__id__in=allowed_facilities
-            )
+        queryset = get_asset_queryset(user=self.request.user, queryset=self.queryset)
         return queryset.annotate(
             latest_status=Subquery(
                 AvailabilityRecord.objects.filter(

diff --git a/care/facility/api/viewsets/bed.py b/care/facility/api/viewsets/bed.py
@@ -30,6 +30,7 @@
 from care.users.models import User
 from care.utils.cache.cache_allowed_facilities import get_accessible_facilities
 from care.utils.filters.choicefilter import CareChoiceFilter, inverse_choices
+from care.utils.queryset.asset_bed import get_asset_bed_queryset, get_bed_queryset
 
 inverse_bed_type = inverse_choices(BedTypeChoices)
 
@@ -76,27 +77,14 @@ class BedViewSet(
     filterset_class = BedFilter
 
     def get_queryset(self):
-        user = self.request.user
-        queryset = self.queryset
-
-        queryset = queryset.annotate(
+        queryset = self.queryset.annotate(
             is_occupied=Exists(
                 ConsultationBed.objects.filter(
                     bed__id=OuterRef("id"), end_date__isnull=True
                 )
             )
         )
-
-        if user.is_superuser:
-            pass
-        elif user.user_type >= User.TYPE_VALUE_MAP["StateLabAdmin"]:
-            queryset = queryset.filter(facility__state=user.state)
-        elif user.user_type >= User.TYPE_VALUE_MAP["DistrictLabAdmin"]:
-            queryset = queryset.filter(facility__district=user.district)
-        else:
-            allowed_facilities = get_accessible_facilities(user)
-            queryset = queryset.filter(facility__id__in=allowed_facilities)
-        return queryset
+        return get_bed_queryset(user=self.request.user, queryset=queryset)
 
     @transaction.atomic
     def create(self, request, *args, **kwargs):
@@ -168,18 +156,7 @@ class AssetBedViewSet(
     lookup_field = "external_id"
 
     def get_queryset(self):
-        user = self.request.user
-        queryset = self.queryset
-        if user.is_superuser:
-            pass
-        elif user.user_type >= User.TYPE_VALUE_MAP["StateLabAdmin"]:
-            queryset = queryset.filter(bed__facility__state=user.state)
-        elif user.user_type >= User.TYPE_VALUE_MAP["DistrictLabAdmin"]:
-            queryset = queryset.filter(bed__facility__district=user.district)
-        else:
-            allowed_facilities = get_accessible_facilities(user)
-            queryset = queryset.filter(bed__facility__id__in=allowed_facilities)
-        return queryset
+        return get_asset_bed_queryset(user=self.request.user, queryset=self.queryset)
 
 
 class PatientAssetBedFilter(filters.FilterSet):
@@ -212,20 +189,9 @@ class PatientAssetBedViewSet(ListModelMixin, GenericViewSet):
     ]
 
     def get_queryset(self):
-        user = self.request.user
-        queryset = self.queryset
-        if user.is_superuser:
-            pass
-        elif user.user_type >= User.TYPE_VALUE_MAP["StateLabAdmin"]:
-            queryset = queryset.filter(bed__facility__state=user.state)
-        elif user.user_type >= User.TYPE_VALUE_MAP["DistrictLabAdmin"]:
-            queryset = queryset.filter(bed__facility__district=user.district)
-        else:
-            allowed_facilities = get_accessible_facilities(user)
-            queryset = queryset.filter(bed__facility__id__in=allowed_facilities)
-        return queryset.filter(
-            bed__facility__external_id=self.kwargs["facility_external_id"]
-        )
+        return get_asset_bed_queryset(
+            user=self.request.user, queryset=self.queryset
+        ).filter(bed__facility__external_id=self.kwargs["facility_external_id"])
 
 
 class ConsultationBedFilter(filters.FilterSet):

diff --git a/care/facility/api/viewsets/camera_preset.py b/care/facility/api/viewsets/camera_preset.py
@@ -5,9 +5,12 @@
 from rest_framework.viewsets import GenericViewSet, ModelViewSet
 
 from care.facility.api.serializers.camera_preset import CameraPresetSerializer
-from care.facility.models import Asset, AssetBed, Bed, CameraPreset
-from care.users.models import User
-from care.utils.cache.cache_allowed_facilities import get_accessible_facilities
+from care.facility.models import CameraPreset
+from care.utils.queryset.asset_bed import (
+    get_asset_bed_queryset,
+    get_asset_queryset,
+    get_bed_queryset,
+)
 
 
 class AssetBedCameraPresetViewSet(ModelViewSet):
@@ -19,19 +22,9 @@ class AssetBedCameraPresetViewSet(ModelViewSet):
     permission_classes = (IsAuthenticated,)
 
     def get_asset_bed_obj(self):
-        user = self.request.user
-        queryset = AssetBed.objects.filter(
+        queryset = get_asset_bed_queryset(self.request.user).filter(
             external_id=self.kwargs["assetbed_external_id"]
         )
-        if user.is_superuser:
-            pass
-        elif user.user_type >= User.TYPE_VALUE_MAP["StateLabAdmin"]:
-            queryset = queryset.filter(bed__facility__state=user.state)
-        elif user.user_type >= User.TYPE_VALUE_MAP["DistrictLabAdmin"]:
-            queryset = queryset.filter(bed__facility__district=user.district)
-        else:
-            allowed_facilities = get_accessible_facilities(user)
-            queryset = queryset.filter(bed__facility__id__in=allowed_facilities)
         return get_object_or_404(queryset)
 
     def get_queryset(self):
@@ -52,35 +45,11 @@ class CameraPresetViewSet(GenericViewSet, ListModelMixin):
     permission_classes = (IsAuthenticated,)
 
     def get_bed_obj(self, external_id: str):
-        user = self.request.user
-        queryset = Bed.objects.filter(external_id=external_id)
-        if user.is_superuser:
-            pass
-        elif user.user_type >= User.TYPE_VALUE_MAP["StateLabAdmin"]:
-            queryset = queryset.filter(facility__state=user.state)
-        elif user.user_type >= User.TYPE_VALUE_MAP["DistrictLabAdmin"]:
-            queryset = queryset.filter(facility__district=user.district)
-        else:
-            allowed_facilities = get_accessible_facilities(user)
-            queryset = queryset.filter(facility__id__in=allowed_facilities)
+        queryset = get_bed_queryset(self.request.user).filter(external_id=external_id)
         return get_object_or_404(queryset)
 
     def get_asset_obj(self, external_id: str):
-        user = self.request.user
-        queryset = Asset.objects.filter(external_id=external_id)
-        if user.is_superuser:
-            pass
-        elif user.user_type >= User.TYPE_VALUE_MAP["StateLabAdmin"]:
-            queryset = queryset.filter(current_location__facility__state=user.state)
-        elif user.user_type >= User.TYPE_VALUE_MAP["DistrictLabAdmin"]:
-            queryset = queryset.filter(
-                current_location__facility__district=user.district
-            )
-        else:
-            allowed_facilities = get_accessible_facilities(user)
-            queryset = queryset.filter(
-                current_location__facility__id__in=allowed_facilities
-            )
+        queryset = get_asset_queryset(self.request.user).filter(external_id=external_id)
         return get_object_or_404(queryset)
 
     def get_queryset(self):

diff --git a/care/utils/queryset/asset_bed.py b/care/utils/queryset/asset_bed.py
@@ -0,0 +1,47 @@
+from care.facility.models import Asset, AssetBed, Bed
+from care.users.models import User
+from care.utils.cache.cache_allowed_facilities import get_accessible_facilities
+
+
+def get_asset_bed_queryset(user, queryset=None):
+    queryset = queryset or AssetBed.objects.all()
+    if user.is_superuser:
+        pass
+    elif user.user_type >= User.TYPE_VALUE_MAP["StateLabAdmin"]:
+        queryset = queryset.filter(bed__facility__state=user.state)
+    elif user.user_type >= User.TYPE_VALUE_MAP["DistrictLabAdmin"]:
+        queryset = queryset.filter(bed__facility__district=user.district)
+    else:
+        allowed_facilities = get_accessible_facilities(user)
+        queryset = queryset.filter(bed__facility__id__in=allowed_facilities)
+    return queryset
+
+
+def get_bed_queryset(user, queryset=None):
+    queryset = queryset or Bed.objects.all()
+    if user.is_superuser:
+        pass
+    elif user.user_type >= User.TYPE_VALUE_MAP["StateLabAdmin"]:
+        queryset = queryset.filter(facility__state=user.state)
+    elif user.user_type >= User.TYPE_VALUE_MAP["DistrictLabAdmin"]:
+        queryset = queryset.filter(facility__district=user.district)
+    else:
+        allowed_facilities = get_accessible_facilities(user)
+        queryset = queryset.filter(facility__id__in=allowed_facilities)
+    return queryset
+
+
+def get_asset_queryset(user, queryset=None):
+    queryset = queryset or Asset.objects.all()
+    if user.is_superuser:
+        pass
+    elif user.user_type >= User.TYPE_VALUE_MAP["StateLabAdmin"]:
+        queryset = queryset.filter(current_location__facility__state=user.state)
+    elif user.user_type >= User.TYPE_VALUE_MAP["DistrictLabAdmin"]:
+        queryset = queryset.filter(current_location__facility__district=user.district)
+    else:
+        allowed_facilities = get_accessible_facilities(user)
+        queryset = queryset.filter(
+            current_location__facility__id__in=allowed_facilities
+        )
+    return queryset