Expand site info metadata (addresses #73)

dataretrieval/nwis.py

@@ -335,8 +335,8 @@ def query_waterservices(service, **kwargs):
     Usage: must specify one major filter: sites, stateCd, bBox,
 
     """
-    if not any(key in kwargs for key in ['sites', 'stateCd', 'bBox', 'huc']):
-        raise TypeError('Query must specify a major filter: sites, stateCd, bBox, or huc')
+    if not any(key in kwargs for key in ['sites', 'stateCd', 'bBox', 'huc', 'countyCd']):
+        raise TypeError('Query must specify a major filter: sites, stateCd, bBox, huc, or countyCd')
 
     if service not in WATERSERVICES_SERVICES:
         raise TypeError('Service not recognized')
@@ -845,12 +845,24 @@ def _set_metadata(response, **parameters):
             unpacked dictionary of the parameters supplied in the request
     """
     md = set_md(response)
-    site_aliases = ['sites', 'site_no']
-    for alias in site_aliases:
-        if alias in parameters:
-            md.site_info = lambda: what_sites(sites=parameters[alias])
-            break
+    # site_no is preferred over sites to set site_info if both are present,
+    # matching behavior of the get_rating() function
+    if 'site_no' in parameters:
+        md.site_info = lambda: what_sites(sites=parameters['site_no'])
+    elif 'sites' in parameters:
+        md.site_info = lambda: what_sites(sites=parameters['sites'])
+    elif 'stateCd' in parameters:
+        md.site_info = lambda: what_sites(stateCd=parameters['stateCd'])
+    elif 'huc' in parameters:
+        md.site_info = lambda: what_sites(huc=parameters['huc'])
+    elif 'countyCd' in parameters:
+        md.site_info = lambda: what_sites(countyCd=parameters['countyCd'])
+    elif 'bBox' in parameters:
+        md.site_info = lambda: what_sites(bBox=parameters['bBox'])
+    else:
+        pass  # don't set metadata site_info attribute
 
+    # define variable_info metadata based on parameterCd if available
     if 'parameterCd' in parameters:
         md.variable_info = lambda: get_pmcodes(parameterCd=parameters['parameterCd'])
 

tests/nwis_test.py

@@ -1,9 +1,13 @@
 import numpy as np
 import pandas as pd
 import pytest
+import requests
 import datetime
 from dataretrieval.nwis import get_record, preformat_peaks_response, get_info
 from dataretrieval.nwis import what_sites, get_iv, get_dv, get_discharge_peaks
+from dataretrieval.nwis import _set_metadata
+import unittest.mock as mock
+
 
 START_DATE = '2018-01-24'
 END_DATE   = '2018-01-25'
@@ -201,3 +205,73 @@ def test_empty_timeseries():
     df = get_record(sites='011277906', service='iv',
                     start='2010-07-20', end='2010-07-20')
     assert df.empty is True
+
+
+class TestMetaData:
+    """Tests of NWIS metadata setting, based on GitHub Issue #73."""
+
+    def test_set_metadata_info_site(self):
+        """Test metadata info is set when site parameter is supplied."""
+        # mock the query response
+        response = mock.MagicMock()
+        # make metadata call
+        md = _set_metadata(response, sites='01491000')
+        # assert that metadata info exists but don't execute lambda function
+        assert md.site_info is not None
+        # assert metadata site_info is callable
+        assert hasattr(md.site_info, '__call__')
+
+    def test_set_metadata_info_site_no(self):
+        """Test metadata info is set when site_no parameter is supplied."""
+        # mock the query response
+        response = mock.MagicMock()
+        # make metadata call
+        md = _set_metadata(response, site_no='01491000')
+        # assert that metadata info exists but don't execute lambda function
+        assert md.site_info is not None
+        # assert metadata site_info is callable
+        assert hasattr(md.site_info, '__call__')
+
+    def test_set_metadata_info_stateCd(self):
+        """Test metadata info is set when stateCd parameter is supplied."""
+        # mock the query response
+        response = mock.MagicMock()
+        # make metadata call
+        md = _set_metadata(response, stateCd='RI')
+        # assert that metadata info exists but don't execute lambda function
+        assert md.site_info is not None
+        # assert metadata site_info is callable
+        assert hasattr(md.site_info, '__call__')
+
+    def test_set_metadata_info_huc(self):
+        """Test metadata info is set when huc parameter is supplied."""
+        # mock the query response
+        response = mock.MagicMock()
+        # make metadata call
+        md = _set_metadata(response, huc='01')
+        # assert that metadata info exists but don't execute lambda function
+        assert md.site_info is not None
+        # assert metadata site_info is callable
+        assert hasattr(md.site_info, '__call__')
+
+    def test_set_metadata_info_bbox(self):
+        """Test metadata info is set when bbox parameter is supplied."""
+        # mock the query response
+        response = mock.MagicMock()
+        # make metadata call
+        md = _set_metadata(response, bBox='-92.8,44.2,-88.9,46.0')
+        # assert that metadata info exists but don't execute lambda function
+        assert md.site_info is not None
+        # assert metadata site_info is callable
+        assert hasattr(md.site_info, '__call__')
+
+    def test_set_metadata_info_countyCd(self):
+        """Test metadata info is set when countyCd parameter is supplied."""
+        # mock the query response
+        response = mock.MagicMock()
+        # make metadata call
+        md = _set_metadata(response, countyCd='01001')
+        # assert that metadata info exists but don't execute lambda function
+        assert md.site_info is not None
+        # assert metadata site_info is callable
+        assert hasattr(md.site_info, '__call__')

tests/waterservices_test.py

@@ -29,7 +29,7 @@ def test_query_waterservices_validation():
     """Tests the validation parameters of the query_waterservices method"""
     with pytest.raises(TypeError) as type_error:
         query_waterservices(service='dv', format='rdb')
-    assert 'Query must specify a major filter: sites, stateCd, bBox, or huc' == str(type_error.value)
+    assert 'Query must specify a major filter: sites, stateCd, bBox, huc, or countyCd' == str(type_error.value)
 
     with pytest.raises(TypeError) as type_error:
         query_waterservices(service=None, sites='sites')
@@ -271,9 +271,7 @@ def assert_metadata(requests_mock, request_url, md, site, parameter_cd, format):
     assert md.url == request_url
     assert isinstance(md.query_time, datetime.timedelta)
     assert md.header == {"mock_header": "value"}
-    if site is None:
-        assert md.site_info is None
-    else:
+    if site is not None:
         site_request_url = "https://waterservices.usgs.gov/nwis/site?sites={}&format=rdb".format(site)
         with open('data/waterservices_site.txt') as text:
             requests_mock.get(site_request_url, text=text.read())