diff --git a/dataretrieval/wqp.py b/dataretrieval/wqp.py index 0b53e387..70d616f9 100644 --- a/dataretrieval/wqp.py +++ b/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) diff --git a/tests/wqp_test.py b/tests/wqp_test.py index f36558bc..a337f7ec 100644 --- a/tests/wqp_test.py +++ b/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"]