fix nan issue when contraband data is missing

nc/tests/api/test_arrests.py

@@ -0,0 +1,45 @@
+from django.urls import reverse
+import pandas as pd
+
+from django.test import TestCase
+import pytest
+
+from nc.models import StopPurpose, DriverRace, DriverEthnicity
+from nc.views.arrests import sort_by_stop_purpose
+from nc.tests.factories import ContrabandFactory, PersonFactory, SearchFactory
+
+
+class ArrestUtilityTests(TestCase):
+    def test_sort_by_stop_purpose(self):
+        """Sort DataFrame by stop_purpose column in order of the IntegerChoices"""
+        df = pd.DataFrame(
+            data={
+                "stop_purpose": [
+                    StopPurpose.CHECKPOINT,
+                    StopPurpose.INVESTIGATION,
+                    StopPurpose.SEAT_BELT_VIOLATION,
+                    StopPurpose.OTHER_MOTOR_VEHICLE_VIOLATION,
+                    StopPurpose.VEHICLE_REGULATORY_VIOLATION,
+                    StopPurpose.VEHICLE_EQUIPMENT_VIOLATION,
+                    StopPurpose.SAFE_MOVEMENT_VIOLATION,
+                    StopPurpose.DRIVING_WHILE_IMPAIRED,
+                    StopPurpose.STOP_LIGHT_SIGN_VIOLATION,
+                    StopPurpose.SPEED_LIMIT_VIOLATION,
+                ],
+            }
+        )
+        self.assertEqual(sort_by_stop_purpose(df)["stop_purpose"].tolist(), StopPurpose.values)
+
+
+@pytest.mark.django_db
+class TestArrests:
+    def test_arrest_contraband_missing_race(self, client, durham):
+        # A single stop will result no data for other races
+        person = PersonFactory(
+            race=DriverRace.BLACK, ethnicity=DriverEthnicity.NON_HISPANIC, stop__agency=durham
+        )
+        search = SearchFactory(stop=person.stop)
+        ContrabandFactory(stop=person.stop, person=person, search=search, pints=2)
+        url = reverse("nc:arrests-percentage-of-stops-per-contraband-type", args=[durham.id])
+        response = client.get(url, data={}, format="json")
+        assert response.status_code == 200

nc/views/arrests.py

@@ -41,8 +41,11 @@ def arrest_table(df, pivot_columns=None, value_key=None, rename_columns=None):
 
 def arrest_query(request, agency_id, group_by, limit_by=None, debug=False):
     """
-    # https://nccopwatch-share.s3.amazonaws.com/2023-10-contraband-type/durham-contraband-hit-rate-type.html
-    # https://nccopwatch-share.s3.amazonaws.com/2024-01-arrest-data/arrest-data-preview-v7.html
+    Query StopSummary view for arrest-related counts.
+
+    Related notebooks:
+    - https://nccopwatch-share.s3.amazonaws.com/2023-10-contraband-type/durham-contraband-hit-rate-type.html
+    - https://nccopwatch-share.s3.amazonaws.com/2024-01-arrest-data/arrest-data-preview-v7.html
     """
     # Build query to filter down queryest
     query = Q(agency_id=agency_id) if agency_id else Q()
@@ -86,10 +89,7 @@ def arrest_query(request, agency_id, group_by, limit_by=None, debug=False):
 
 
 def contraband_query(request, agency_id, group_by, limit_by=None, debug=False):
-    """
-    # https://nccopwatch-share.s3.amazonaws.com/2023-10-contraband-type/durham-contraband-hit-rate-type.html
-    # https://nccopwatch-share.s3.amazonaws.com/2024-01-arrest-data/arrest-data-preview-v7.html
-    """
+    """Query ContrabandSummary for arrest-related counts."""
     # Build query to filter down queryest
     query = Q(agency_id=agency_id) if agency_id else Q()
     officer_id = request.query_params.get("officer", None)
@@ -128,6 +128,7 @@ def contraband_query(request, agency_id, group_by, limit_by=None, debug=False):
         df.contraband_and_driver_arrest_count / df.contraband_count
     )
     df["driver_stop_arrest_rate"] = df.contraband_and_driver_arrest_count / df.stop_count
+    df.fillna(0, inplace=True)
     if "driver_race_comb" in group_by:
         # Add custom sortable driver race column
         columns = ["White", "Black", "Hispanic", "Asian", "Native American", "Other"]