Merge remote-tracking branch 'origin/main' into fix/nldi-cleanup

# Conflicts:
#	dataretrieval/nldi.py

Author: thodson-usgs

dataretrieval/waterdata/utils.py

@@ -892,11 +892,13 @@ def _handle_stats_nesting(
     # otherwise return a geodataframe
     if not geopd:
         df = pd.json_normalize(body["features"]).drop(
-            columns=["type", "properties.data"]
+            columns=["type", "properties.data"], errors="ignore"
         )
         df.columns = df.columns.str.split(".").str[-1]
     else:
-        df = gpd.GeoDataFrame.from_features(body["features"]).drop(columns=["data"])
+        df = gpd.GeoDataFrame.from_features(body["features"]).drop(
+            columns=["data"], errors="ignore"
+        )
 
     # Unnest json features, properties, data, and values while retaining necessary
     # metadata to merge with main dataframe.

dataretrieval/wqp.py

@@ -127,31 +127,23 @@ def get_results(
     kwargs = _check_kwargs(kwargs)
 
     if legacy is True:
-        if (
-            "dataProfile" in kwargs
-            and kwargs["dataProfile"] not in result_profiles_legacy
-        ):
-            raise TypeError(
-                f"dataProfile {kwargs['dataProfile']} is not a legacy profile.",
-                f"Valid options are {result_profiles_legacy}.",
-            )
-
+        valid_profiles = result_profiles_legacy
+        kind = "legacy"
         url = wqp_url("Result")
-
     else:
-        if (
-            "dataProfile" in kwargs
-            and kwargs["dataProfile"] not in result_profiles_wqx3
-        ):
-            raise TypeError(
-                f"dataProfile {kwargs['dataProfile']} is not a valid WQX3.0"
-                f"profile. Valid options are {result_profiles_wqx3}.",
-            )
-        else:
-            kwargs["dataProfile"] = "fullPhysChem"
-
+        valid_profiles = result_profiles_wqx3
+        kind = "WQX3.0"
         url = wqx3_url("Result")
 
+    profile = kwargs.get("dataProfile")
+    if profile is not None and profile not in valid_profiles:
+        raise ValueError(
+            f"dataProfile {profile!r} is not a valid {kind} profile. "
+            f"Valid options are {valid_profiles}."
+        )
+    if legacy is not True and profile is None:
+        kwargs["dataProfile"] = "fullPhysChem"
+
     response = query(url, kwargs, delimiter=";", ssl_check=ssl_check)
 
     df = pd.read_csv(StringIO(response.text), delimiter=",", low_memory=False)

tests/waterdata_utils_test.py

@@ -4,6 +4,7 @@
 
 from dataretrieval.waterdata.utils import (
     _get_args,
+    _handle_stats_nesting,
     _walk_pages,
 )
 
@@ -80,3 +81,33 @@ def test_walk_pages_multiple_mocked():
     assert mock_client.send.called
     assert mock_client.request.called
     assert mock_client.request.call_args[0][1] == "https://example.com/page2"
+
+
+def test_handle_stats_nesting_tolerates_missing_drop_columns():
+    """If the upstream stats response shape ever changes such that one of
+    the columns we try to drop ("type", "properties.data") is absent, the
+    function should still return a DataFrame instead of raising KeyError.
+    """
+    body = {
+        "next": None,
+        "features": [
+            {
+                "properties": {
+                    "monitoring_location_id": "USGS-12345",
+                    "data": [
+                        {
+                            "parameter_code": "00060",
+                            "unit_of_measure": "ft^3/s",
+                            "parent_time_series_id": "ts-1",
+                            "values": [{"statistic_id": "mean", "value": 10.0}],
+                        }
+                    ],
+                },
+            }
+        ],
+    }
+
+    df = _handle_stats_nesting(body, geopd=False)
+
+    assert len(df) == 1
+    assert df["monitoring_location_id"].iloc[0] == "USGS-12345"

tests/wqp_test.py

@@ -216,3 +216,25 @@ def test_check_kwargs():
     kwargs = {"mimeType": "foo"}
     with pytest.raises(ValueError):
         kwargs = _check_kwargs(kwargs)
+
+
+def test_get_results_wqx3_preserves_user_dataProfile(requests_mock):
+    """A valid user-supplied WQX3.0 profile must not be overwritten.
+
+    Regression: previously the `else` branch of the `dataProfile` validation
+    triggered whenever the value was *not invalid*, including any valid
+    user-supplied profile, silently overwriting it with 'fullPhysChem'.
+    """
+    request_url = (
+        "https://www.waterqualitydata.us/wqx3/Result/search?"
+        "siteid=UTAHDWQ_WQX-4993795&mimeType=csv&dataProfile=narrow"
+    )
+    response_file_path = "tests/data/wqp3_results.txt"
+    mock_request(requests_mock, request_url, response_file_path)
+
+    df, _md = get_results(
+        legacy=False, siteid="UTAHDWQ_WQX-4993795", dataProfile="narrow"
+    )
+    assert isinstance(df, DataFrame)
+    sent = requests_mock.request_history[-1]
+    assert sent.qs.get("dataprofile") == ["narrow"]