Implemented sanity check on epoch (#2)

AdvancedClimateSystems · Oct 31, 2016 · ce0cec2 · ce0cec2
1 parent 420a8a8
commit ce0cec2
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 1 deletion.
diff --git a/inflow/connection.py b/inflow/connection.py
@@ -106,14 +106,18 @@ def get_write_url(self):
         return url
 
     def get_query_url(self, query, epoch):
-        """ Returns the url needed to query measurements from InfluxDB. """
+        """ Returns the url needed to query measurements from InfluxDB.  """
         url = '{}/query?db={}&q={}'.format(
             self.uri,
             self.db,
             quote_plus(query)
         )
 
         if epoch is not None:
+            if epoch not in ['h', 'm', 's', 'ms', 'u', 'ns']:
+                raise ValueError('Invalid epoch provided, must be one of '
+                                 'h, m, s, ms, u or ns.')
+
             url = '{}&epoch={}'.format(url, epoch)
 
         return url

diff --git a/tests/test_client.py b/tests/test_client.py
@@ -375,6 +375,10 @@ def test_should_use_epoch_on_query(self, client, get):
         url = get.call_args[0][0]
         assert 'epoch=s' in url
 
+    def test_should_throw_on_invalid_epoch(self, client, get):
+        with pytest.raises(ValueError):
+            client.query('SELECT * FROM "temperature"', epoch='not an epoch')
+
     def test_should_quote_query(self, client, get):
         client.query('SELECT * FROM "temperatures"')