Fix batch writing for influxdb08 client.

geodimm · geodimm · commit a74224d50969 · 2015-04-25T01:57:50.000+01:00
This fixes issue influxdata#102
diff --git a/influxdb/influxdb08/client.py b/influxdb/influxdb08/client.py
@@ -291,11 +291,11 @@ def list_chunks(l, n):
                 yield l[i:i + n]
 
         batch_size = kwargs.get('batch_size')
-        if batch_size:
+        if batch_size and batch_size > 0:
             for item in data:
                 name = item.get('name')
                 columns = item.get('columns')
-                point_list = item.get('points')
+                point_list = item.get('points', [])
 
                 for batch in list_chunks(point_list, batch_size):
                     item = [{
@@ -306,10 +306,10 @@ def list_chunks(l, n):
                     self._write_points(
                         data=item,
                         time_precision=time_precision)
-
-                return True
-
-        return self._write_points(data=data, time_precision=time_precision)
+            return True
+        else:
+            return self._write_points(data=data,
+                                      time_precision=time_precision)
 
     def write_points_with_precision(self, data, time_precision='s'):
         """
diff --git a/tests/influxdb/influxdb08/client_test.py b/tests/influxdb/influxdb08/client_test.py
@@ -186,12 +186,46 @@ def test_write_points_string(self):
             )
 
     def test_write_points_batch(self):
-        with _mocked_session('post', 200, self.dummy_points):
-            cli = InfluxDBClient('host', 8086, 'username', 'password', 'db')
-            assert cli.write_points(
-                data=self.dummy_points,
-                batch_size=2
-            ) is True
+        with requests_mock.Mocker() as m:
+            m.register_uri(requests_mock.POST,
+                           "http://localhost:8086/db/db/series")
+            cli = InfluxDBClient('localhost', 8086,
+                                 'username', 'password', 'db')
+            cli.write_points(data=self.dummy_points, batch_size=2)
+        self.assertEqual(1, m.call_count)
+
+    def test_write_points_batch_invalid_size(self):
+        with requests_mock.Mocker() as m:
+            m.register_uri(requests_mock.POST,
+                           "http://localhost:8086/db/db/series")
+            cli = InfluxDBClient('localhost', 8086,
+                                 'username', 'password', 'db')
+            cli.write_points(data=self.dummy_points, batch_size=-2)
+        self.assertEqual(1, m.call_count)
+
+    def test_write_points_batch_multiple_series(self):
+        dummy_points = [
+            {"points": [["1", 1, 1.0], ["2", 2, 2.0], ["3", 3, 3.0],
+                        ["4", 4, 4.0], ["5", 5, 5.0]],
+             "name": "foo",
+             "columns": ["val1", "val2", "val3"]},
+            {"points": [["1", 1, 1.0], ["2", 2, 2.0], ["3", 3, 3.0],
+                        ["4", 4, 4.0], ["5", 5, 5.0], ["6", 6, 6.0],
+                        ["7", 7, 7.0], ["8", 8, 8.0]],
+             "name": "bar",
+             "columns": ["val1", "val2", "val3"]},
+        ]
+        expected_last_body = [{'points': [['7', 7, 7.0], ['8', 8, 8.0]],
+                               'name': 'bar',
+                               'columns': ['val1', 'val2', 'val3']}]
+        with requests_mock.Mocker() as m:
+            m.register_uri(requests_mock.POST,
+                           "http://localhost:8086/db/db/series")
+            cli = InfluxDBClient('localhost', 8086,
+                                 'username', 'password', 'db')
+            cli.write_points(data=dummy_points, batch_size=3)
+        self.assertEqual(m.call_count, 5)
+        self.assertEqual(expected_last_body, m.request_history[4].json())
 
     def test_write_points_udp(self):
         s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
