[DEV-15452] Mythos finding - district_original location filter silently overwrites district_current

usaspending_api/awards/tests/unit/test_location_filter_geocode.py

@@ -2,14 +2,15 @@
 
 from usaspending_api.awards.v2.filters.location_filter_geocode import (
     create_nested_object,
+    geocode_filter_locations,
     get_fields_list,
     location_error_handling,
     validate_location_keys,
 )
 from usaspending_api.common.exceptions import InvalidParameterException
 from usaspending_api.common.helpers.api_helper import (
-    INCOMPATIBLE_DISTRICT_LOCATION_PARAMETERS,
     DUPLICATE_DISTRICT_LOCATION_PARAMETERS,
+    INCOMPATIBLE_DISTRICT_LOCATION_PARAMETERS,
 )
 
 
@@ -101,3 +102,24 @@ def test_get_fields_list():
     assert get_fields_list("county_code", "01") == ["1", "01", "1.0"]
     assert get_fields_list("feet", "01") == ["01"]
     assert get_fields_list("congressional_code", "abc") == ["abc"]
+
+
+def test_geocode_filter_locations_for_both_district_original_and_current_across_elements():
+    """
+    Cross-element same-state current and original must OR not overwrite
+    """
+
+    values = [
+        {"country": "USA", "state": "VA", "district_original": "08"},
+        {"country": "USA", "state": "VA", "district_current": "11"}
+    ]
+    # validate_location_keys only checks for moth district types within a single element
+    assert validate_location_keys(values) is None
+    nested = create_nested_object(values)
+    assert nested["USA"]["VA"]["district_original"] == ["08"]
+    assert nested["USA"]["VA"]["district_current"] == ["11"]
+
+    q = geocode_filter_locations("test", values)
+    q_repr = repr(q)
+    assert "test_congressional_code_current__in" in q_repr
+    assert "test_congressional_code__in" in q_repr

usaspending_api/awards/v2/filters/location_filter_geocode.py

@@ -1,3 +1,5 @@
+from typing import Any
+
 from django.db.models import Q
 
 from usaspending_api.common.exceptions import InvalidParameterException
@@ -10,57 +12,72 @@
 ALL_FOREIGN_COUNTRIES = "FOREIGN"
 
 
-def geocode_filter_locations(scope: str, values: list) -> Q:
+def geocode_filter_locations(scope: str, values: list[dict[str, Any]]) -> Q:
     """
     Function filter querysets on location table
     scope- place of performance or recipient location mappings
     values- array of location requests
     returns queryset
     """
-    or_queryset = Q()
-
-    # creates a dictionary with all of the locations organized by country
-    # Counties and congressional districts are nested under state codes
     nested_values = create_nested_object(values)
 
-    # In this for-loop a django Q filter object is created from the python dict
+    or_queryset = Q()
     for country, state_zip in nested_values.items():
-        country_qs = None
-        if country != ALL_FOREIGN_COUNTRIES:
-            country_qs = Q(**{f"{scope}_country_code__exact": country})
-        state_qs = Q()
-
-        for state_zip_key, location_values in state_zip.items():
-
-            if state_zip_key == "city":
-                state_inner_qs = Q(**{f"{scope}_city_name__in": location_values})
-            elif state_zip_key == "zip":
-                state_inner_qs = Q(**{f"{scope}_zip5__in": location_values})
-            else:
-                state_inner_qs = Q(**{f"{scope}_state_code__exact": state_zip_key.upper()})
-                county_qs = Q()
-                district_qs = Q()
-                city_qs = Q()
-
-                if location_values["county"]:
-                    county_qs = Q(**{f"{scope}_county_code__in": location_values["county"]})
-                if location_values["district_current"]:
-                    district_qs = Q(**{f"{scope}_congressional_code_current__in": location_values["district_current"]})
-                if location_values["district_original"]:
-                    district_qs = Q(**{f"{scope}_congressional_code__in": location_values["district_original"]})
-                if location_values["city"]:
-                    city_qs = Q(**{f"{scope}_city_name__in": location_values["city"]})
-                state_inner_qs &= county_qs | district_qs | city_qs
-
-            state_qs |= state_inner_qs
-        if country_qs:
-            or_queryset |= country_qs & state_qs
-        else:
-            or_queryset |= state_qs
+        country_qs = _build_country_query(scope, country)
+        state_qs = _build_state_queries(scope, state_zip)
+
+        or_queryset |= (country_qs & state_qs) if country_qs else state_qs
+
     return or_queryset
 
 
-def validate_location_keys(values):
+def _build_country_query(scope: str, country: str) -> Q | None:
+    """Build country-level query filter"""
+    if country == ALL_FOREIGN_COUNTRIES:
+        return None
+    return Q(**{f"{scope}_country_code__exact": country})
+
+
+def _build_state_queries(scope: str, state_zip: dict[str, Any]) -> Q:
+    """Build state-level query filters"""
+    state_qs = Q()
+
+    for state_zip_key, location_values in state_zip.items():
+        if state_zip_key == "city":
+            state_inner_qs = Q(**{f"{scope}_city_name__in": location_values})
+        elif state_zip_key == "zip":
+            state_inner_qs = Q(**{f"{scope}_zip5__in": location_values})
+        else:
+            state_inner_qs = _build_state_location_query(scope, state_zip_key, location_values)
+
+        state_qs |= state_inner_qs
+
+    return state_qs
+
+
+def _build_state_location_query(scope: str, state_code: str, location_values: dict[str, Any]) -> Q:
+    """Build query for state with nested county/district/city filters"""
+    state_qs = Q(**{f"{scope}_state_code__exact": state_code.upper()})
+
+    # Build sub-filters
+    sub_filters = Q()
+
+    if location_values.get("county"):
+        sub_filters |= Q(**{f"{scope}_county_code__in": location_values["county"]})
+
+    if location_values.get("district_current"):
+        sub_filters |= Q(**{f"{scope}_congressional_code_current__in": location_values["district_current"]})
+
+    if location_values.get("district_original"):
+        sub_filters |= Q(**{f"{scope}_congressional_code__in": location_values["district_original"]})
+
+    if location_values.get("city"):
+        sub_filters |= Q(**{f"{scope}_city_name__in": location_values["city"]})
+
+    return state_qs & sub_filters
+
+
+def validate_location_keys(values: list[dict[str, Any]]) -> None:
     """Validate that the keys provided are sufficient and match properly."""
     for v in values:
         state = v.get("state")
@@ -78,70 +95,107 @@ def validate_location_keys(values):
             location_error_handling(v.keys())
 
 
-def create_nested_object(values):
+def create_nested_object(values: list[dict[str, Any]]) -> dict[str, Any]:
     """Makes sure keys provided are valid"""
     validate_location_keys(values)
 
-    nested_locations = {}
+    nested_locations: dict[str, Any] = {}
     for v in values:
         upper_case_dict_values(v)
-        city = v.get("city")
-        country = v.get("country")
-        county = v.get("county")
-        district_original = v.get("district_original")
-        district_current = v.get("district_current")
-        state = v.get("state")
-        zip = v.get("zip")
-        # First level in location filtering in country
-        # All location requests must have a country otherwise there will be a key error
-        if nested_locations.get(country) is None:
-            nested_locations[country] = {}
-
-        # Initialize the list
-        if zip and not nested_locations[country].get("zip"):
-            nested_locations[country]["zip"] = []
-
-        if city and not nested_locations[country].get("city"):
-            nested_locations[country]["city"] = []
-
-        # Second level of filtering is zip and state
-        # Requests must have a country+zip or country+state combination
-        if zip:
-            # Appending zips so we don't overwrite
-            nested_locations[country]["zip"].append(zip)
-
-        # If we have a state, add it to the list
-        if state and nested_locations[country].get(state) is None:
-            nested_locations[country][state] = {
-                "county": [],
-                "district_original": [],
-                "district_current": [],
-                "city": [],
-            }
-
-        # Based on previous checks, there will always be a state if either of these exist
-        if county:
-            nested_locations[country][state]["county"].extend(get_fields_list("county", county))
-
-        if district_current:
-            nested_locations[country][state]["district_current"].extend(
-                get_fields_list("district_current", district_current)
-            )
-
-        if district_original:
-            nested_locations[country][state]["district_original"].extend(
-                get_fields_list("district_original", district_original)
-            )
-
-        if city and state:
-            nested_locations[country][state]["city"].append(city)
-        elif city:
-            nested_locations[country]["city"].append(city)
+        location_data = _extract_location_data(v)
+        _process_location(nested_locations, location_data)
 
     return nested_locations
 
 
-def location_error_handling(fields):
+def _extract_location_data(location_dict: dict[str, Any]) -> dict[str, Any]:
+    """Extract location fields from the input dictionary"""
+    return {
+        "city": location_dict.get("city"),
+        "country": location_dict.get("country"),
+        "county": location_dict.get("county"),
+        "district_original": location_dict.get("district_original"),
+        "district_current": location_dict.get("district_current"),
+        "state": location_dict.get("state"),
+        "zip_code": location_dict.get("zip"),
+    }
+
+
+def _process_location(nested_locations: dict[str, Any], location_data: dict[str, Any]) -> None:
+    """Process a single location and add it to the nested structure"""
+    country = location_data["country"]
+
+    # Initialize country if needed
+    if country not in nested_locations:
+        nested_locations[country] = {}
+
+    # Process zip codes
+    _process_zip_code(nested_locations[country], location_data["zip_code"])
+
+    # Process city (country-level)
+    _process_country_level_city(nested_locations[country], location_data["city"])
+
+    # Process state-level data
+    if location_data["state"]:
+        _process_state_data(nested_locations[country], location_data)
+    elif location_data["city"]:
+        # City without state goes to country level
+        nested_locations[country]["city"].append(location_data["city"])
+
+
+def _process_zip_code(country_data: dict[str, Any], zip_code: str | None) -> None:
+    """Process zip code for a country"""
+    if not zip_code:
+        return
+
+    if "zip" not in country_data:
+        country_data["zip"] = []
+
+    country_data["zip"].append(zip_code)
+
+
+def _process_country_level_city(country_data: dict[str, Any], city: str | None) -> None:
+    """Initialize city list at country level if needed"""
+    if city and "city" not in country_data:
+        country_data["city"] = []
+
+
+def _process_state_data(country_data: dict[str, Any], location_data: dict[str, Any]) -> None:
+    """Process state-level location data"""
+    state = location_data["state"]
+
+    # Initialize state structure if needed
+    if state not in country_data:
+        country_data[state] = {
+            "county": [],
+            "district_original": [],
+            "district_current": [],
+            "city": [],
+        }
+
+    state_data = country_data[state]
+
+    # Add county data
+    if location_data["county"]:
+        state_data["county"].extend(get_fields_list("county", location_data["county"]))
+
+    # Add district data
+    if location_data["district_current"]:
+        state_data["district_current"].extend(
+            get_fields_list("district_current", location_data["district_current"])
+        )
+
+    if location_data["district_original"]:
+        state_data["district_original"].extend(
+            get_fields_list("district_original", location_data["district_original"])
+        )
+
+    # Add city data
+    if location_data["city"]:
+        state_data["city"].append(location_data["city"])
+
+
+def location_error_handling(fields: Any) -> None:
     """Raise the relevant error for location keys."""
     # Request must have country, and can only have 3 fields, and must have state if there is county
     if "country" not in fields:
@@ -151,7 +205,7 @@ def location_error_handling(fields):
         raise InvalidParameterException("Invalid filter:  Missing necessary location field: state.")
 
 
-def get_fields_list(scope, field_value):
+def get_fields_list(scope: str, field_value: str) -> list[str]:
     """List of values to search for; `field_value`, plus possibly variants on it"""
     if scope in ["congressional_code", "county_code"]:
         try: