import datetime
import ulmo
import test_util
import numpy as np

def test_get_stations():
    stations_file = 'cdec/historical/all_stations.csv'
    with test_util.mocked_urls(stations_file):
        stations = ulmo.cdec.historical.get_stations()
    assert 2000 < len(stations)
    assert u'PRA' in stations.index


def test_get_sensors():
    sensors_file = 'cdec/historical/sensors.htm'
    with test_util.mocked_urls(sensors_file):
        sensors = ulmo.cdec.historical.get_sensors()
    assert 200 < len(sensors)
    assert u'FLOW, RIVER DISCHARGE' == sensors.ix[20]['Description']


def test_get_station_sensors():
    test_sets = [
            ('PRA', [6,6,15,15,22,23,76])
    ]

    for station_id, test_value in test_sets:

        stations_file = 'cdec/historical/%s.htm' % (station_id)
        with test_util.mocked_urls(stations_file):
            available_sensors = ulmo.cdec.historical.get_station_sensors([station_id])

        assert available_sensors[station_id].sensor_id.to_dict().values() == test_value


def test_get_station_data():
    test_sets = [
            (u'PRA', u'RESERVOIR ELEVATION daily', ['2000-01-01', '2000-01-02'])
    ]

    for station_id, var_name, test_values in test_sets:
        data_regex = 'http://cdec.water.ca.gov/cgi-progs/queryCSV?station_id.*'
        data_file = 'cdec/historical/%s.csv' % station_id
        sensors_regex = 'http://cdec.water.ca.gov/misc/senslist.html'
        sensors_file = 'cdec/historical/sensors.htm'

        url_files = {sensors_regex: sensors_file, data_regex: data_file}
        with test_util.mocked_urls(url_files):
            station_data = ulmo.cdec.historical.get_data(['PRA'], [6],
                resolutions=['daily'], start='2000-1-1', end='2000-1-2')

            assert np.all(test_values == station_data[station_id][var_name].index)

@@ -4,16 +4,40 @@
 ort 
-datetime
+pandas as pd%0Aimport numpy as np%0A
 %0Aimp
@@ -62,34 +62,16 @@
 st_util%0A
-import numpy as np
 %0A%0Adef te
@@ -518,16 +518,69 @@
 ensors)%0A
+    sensors.index = sensors%5B'Sensor No'%5D.astype(int)%0A
     asse
@@ -622,14 +622,14 @@
 ors.
-ix%5B20%5D
+xs(20)
 %5B'De
@@ -1798,16 +1798,85 @@
 -1-2')%0A%0A
+            test_timestamps = %5Bpd.Timestamp(t) for t in test_values%5D%0A
         
@@ -1898,21 +1898,25 @@
 ll(test_
-value
+timestamp
 s == sta

fix cdec historical test to be more general with how it tests values these comparisons start failing with pandas 0.13