added more tests for instantaneous values

lauraschachter · lauraschachter · commit 9c0527bdeaa9 · 2020-10-28T14:09:05.000-05:00
diff --git a/pydrograph/nwis.py b/pydrograph/nwis.py
@@ -535,20 +535,29 @@ def get_ivs(self, station_ID, parameter_code='00060', start_date='2000-01-01', e
         cols = sitefile_text[skiprows - 2].decode('utf-8').strip().split('\t')
         loginfo = [str(station_ID), url, self.get_datetime_retrieved(sitefile_text)]
         df = pd.read_csv(url, sep='\t', skiprows=skiprows, header=None, names=cols, dtype={'site_no': object})
-        if len(df) > 0:
-            df.index = pd.to_datetime(df.datetime)
-            loginfo.append(True)
-        else:
-            loginfo.append(False)
-        self.log = self.log.append(pd.DataFrame([loginfo], columns=self.log_cols))
+
+        if len(df) > 2:
+
+            if len(df) > 0:
+                df.index = pd.to_datetime(df.datetime)
+                loginfo.append(True)
+            else:
+                loginfo.append(False)
+            self.log = self.log.append(pd.DataFrame([loginfo], columns=self.log_cols))
         
-        if sample_period is not None:
-            df = df.resample(sample_period).agg(agg_method)
-            df = df.rename(columns = {df.columns[1]: 'discharge (cfs)'})    
+            if sample_period is not None:
+                df = df.resample(sample_period).agg(agg_method)
+                df = df.rename(columns = {df.columns[0]: 'discharge (cfs)'})
+            else:
+                df = df.rename(columns = {df.columns[4]: 'discharge (cfs)', df.columns[5]: 'code'})
+
         else:
-            df = df.rename(columns = {df.columns[4]: 'discharge (cfs)', df.columns[5]: 'code'})
-        
-        return df
+            print('No data at this site during this timeframe.')
+
+        if len(df) > 2:
+            return df
+        else:
+            return None
 
     def get_measurements(self, station_ID, txt='measurement'):
         """Retrieves field measurements for a site.
diff --git a/pydrograph/tests/test_nwis.py b/pydrograph/tests/test_nwis.py
@@ -7,25 +7,28 @@
 import numpy as np
 from pandas._testing import assert_frame_equal
 
-
 @pytest.fixture(scope='session')
 def extent_poly():
     extent_poly = box(390000, 1330000, 500000, 1455000)
     extent_poly_ll = project(extent_poly, "+init=epsg:{}".format(5070), "+init=epsg:4269")
     return extent_poly_ll
 
-
 @pytest.fixture(scope='session')
 def nwis_instance(extent_poly):
     nwis_instance = Nwis(extent=extent_poly)
     return nwis_instance
 
-
 @pytest.fixture(scope='session')
 def field_sites(nwis_instance):
     field_sites = nwis_instance.get_siteinfo('field_measurements')
     return field_sites
 
+@pytest.fixture(scope='session')
+def stations(nwis_instance):
+    stations = nwis_instance.get_iv_siteinfo(attributes='iv_attributes')
+    stations = stations.site_no.unique()
+    #stations = stations.to_list()
+    return stations
 
 def test_compute_geometries(extent_poly, nwis_instance, field_sites):
     geoms = nwis_instance._compute_geometries(field_sites)
@@ -38,3 +41,10 @@ def test_instantaneous_value(nwis_instance):
     assert len(df) > 0
     assert 'site_no' in df.columns
     assert df.site_no.dtype == np.object
+
+def test_get_all_ivs(nwis_instance, stations):
+    all_sites = nwis_instance.get_all_ivs(stations)
+    site_one = list(all_sites.values())[0]
+    assert len(all_sites) > 0
+    assert len(site_one) > 2
+    #assert all_sites.site_no.dtype == np.object
