Added test_auto_commit

aviau · aviau · commit d78df7601bb9 · 2015-01-24T13:00:29.000-05:00
diff --git a/influxdb/helper.py b/influxdb/helper.py
@@ -108,7 +108,9 @@ def __init__(self, **kw):
 
         cls._datapoints[cls._series_name.format(**kw)].append(cls._type(**kw))
 
-        if cls._autocommit and len(cls._datapoints) >= cls._bulk_size:
+        if cls._autocommit and \
+                sum(len(series) for series in cls._datapoints.values()) \
+                >= cls._bulk_size:
             cls.commit()
 
     @classmethod
diff --git a/tests/influxdb/helper_test.py b/tests/influxdb/helper_test.py
@@ -3,6 +3,7 @@
 import unittest
 import warnings
 
+import mock
 from influxdb import SeriesHelper, InfluxDBClient
 from requests.exceptions import ConnectionError
 
@@ -18,7 +19,8 @@ def setUpClass(cls):
             8086,
             'username',
             'password',
-            'database')
+            'database'
+        )
 
         class MySeriesHelper(SeriesHelper):
 
@@ -31,10 +33,33 @@ class Meta:
 
         TestSeriesHelper.MySeriesHelper = MySeriesHelper
 
+    def test_auto_commit(self):
+        """
+        Tests that write_points is called after the right number of events
+        """
+        class AutoCommitTest(SeriesHelper):
+
+            class Meta:
+                series_name = 'events.stats.{server_name}'
+                fields = ['time', 'server_name']
+                bulk_size = 5
+                client = InfluxDBClient()
+                autocommit = True
+
+        fake_write_points = mock.MagicMock()
+        AutoCommitTest(server_name='us.east-1', time=159)
+        AutoCommitTest._client.write_points = fake_write_points
+        AutoCommitTest(server_name='us.east-1', time=158)
+        AutoCommitTest(server_name='us.east-1', time=157)
+        AutoCommitTest(server_name='us.east-1', time=156)
+        self.assertFalse(fake_write_points.called)
+        AutoCommitTest(server_name='us.east-1', time=3443)
+        self.assertTrue(fake_write_points.called)
+
     def testSingleSeriesName(self):
-        '''
+        """
         Tests JSON conversion when there is only one series name.
-        '''
+        """
         TestSeriesHelper.MySeriesHelper(server_name='us.east-1', time=159)
         TestSeriesHelper.MySeriesHelper(server_name='us.east-1', time=158)
         TestSeriesHelper.MySeriesHelper(server_name='us.east-1', time=157)
@@ -121,9 +146,9 @@ class Meta:
                                         'server_name': 'us.east-1'})
 
     def testWarnBulkSizeZero(self):
-        '''
+        """
         Tests warning for an invalid bulk size.
-        '''
+        """
         class WarnBulkSizeZero(SeriesHelper):
 
             class Meta:
@@ -148,9 +173,9 @@ class Meta:
                           'Warning message did not contain "forced to 1".')
 
     def testWarnBulkSizeNoEffect(self):
-        '''
+        """
         Tests warning for a set bulk size but autocommit False.
-        '''
+        """
         class WarnBulkSizeNoEffect(SeriesHelper):
 
             class Meta:
