Repository: mixpanel/mixpanel-python

Issue:
flush function for Buffered Consumer not working
Hi,
in class BufferedConsumer the flush function in line 338 should change to 
def flush (self,api_key=None) 

and then in line 444-445 should change to:
        for endpoint in self._buffers.keys():
            self._flush_endpoint(endpoint,api_key=api_key)


Discussion:

Patch:
diff --git a/mixpanel/__init__.py b/mixpanel/__init__.py
index fc17602..e8ab564 100644
--- a/mixpanel/__init__.py
+++ b/mixpanel/__init__.py
@@ -345,6 +345,7 @@ class Consumer(object):
         :param endpoint: the Mixpanel API endpoint appropriate for the message
         :type endpoint: "events" | "people" | "imports"
         :param str json_message: a JSON message formatted for the endpoint
+        :param str api_key: your Mixpanel project's API key
         :raises MixpanelException: if the endpoint doesn't exist, the server is
             unreachable, or the message cannot be processed
         """
@@ -412,6 +413,7 @@ class BufferedConsumer(object):
             'imports': [],
         }
         self._max_size = min(50, max_size)
+        self._api_key = None
 
     def send(self, endpoint, json_message, api_key=None):
         """Record an event or profile update.
@@ -424,16 +426,22 @@ class BufferedConsumer(object):
         :param endpoint: the Mixpanel API endpoint appropriate for the message
         :type endpoint: "events" | "people" | "imports"
         :param str json_message: a JSON message formatted for the endpoint
+        :param str api_key: your Mixpanel project's API key
         :raises MixpanelException: if the endpoint doesn't exist, the server is
             unreachable, or any buffered message cannot be processed
+
+        .. versionadded:: 4.3.2
+            The *api_key* parameter.
         """
         if endpoint not in self._buffers:
             raise MixpanelException('No such endpoint "{0}". Valid endpoints are one of {1}'.format(endpoint, self._buffers.keys()))
 
         buf = self._buffers[endpoint]
         buf.append(json_message)
+        if api_key is not None:
+            self._api_key = api_key
         if len(buf) >= self._max_size:
-            self._flush_endpoint(endpoint, api_key)
+            self._flush_endpoint(endpoint)
 
     def flush(self):
         """Immediately send all buffered messages to Mixpanel.
@@ -444,13 +452,13 @@ class BufferedConsumer(object):
         for endpoint in self._buffers.keys():
             self._flush_endpoint(endpoint)
 
-    def _flush_endpoint(self, endpoint, api_key=None):
+    def _flush_endpoint(self, endpoint):
         buf = self._buffers[endpoint]
         while buf:
             batch = buf[:self._max_size]
             batch_json = '[{0}]'.format(','.join(batch))
             try:
-                self._consumer.send(endpoint, batch_json, api_key)
+                self._consumer.send(endpoint, batch_json, self._api_key)
             except MixpanelException as orig_e:
                 mp_e = MixpanelException(orig_e)
                 mp_e.message = batch_json


Test patch:
diff --git a/test_mixpanel.py b/test_mixpanel.py
index 132b2f6..190517c 100644
--- a/test_mixpanel.py
+++ b/test_mixpanel.py
@@ -353,40 +353,32 @@ class TestBufferedConsumer:
     def setup_class(cls):
         cls.MAX_LENGTH = 10
         cls.consumer = mixpanel.BufferedConsumer(cls.MAX_LENGTH)
-        cls.mock = Mock()
-        cls.mock.read.return_value = six.b('{"status":1, "error": null}')
+        cls.consumer._consumer = LogConsumer()
+        cls.log = cls.consumer._consumer.log
 
-    def test_buffer_hold_and_flush(self):
-        with patch('six.moves.urllib.request.urlopen', return_value=self.mock) as urlopen:
-            self.consumer.send('events', '"Event"')
-            assert not self.mock.called
-            self.consumer.flush()
+    def setup_method(self):
+        del self.log[:]
 
-            assert urlopen.call_count == 1
-
-            (call_args, kwargs) = urlopen.call_args
-            (request,) = call_args
-            timeout = kwargs.get('timeout', None)
-
-            assert request.get_full_url() == 'https://api.mixpanel.com/track'
-            assert qs(request.data) == qs('ip=0&data=WyJFdmVudCJd&verbose=1')
-            assert timeout is None
+    def test_buffer_hold_and_flush(self):
+        self.consumer.send('events', '"Event"')
+        assert len(self.log) == 0
+        self.consumer.flush()
+        assert self.log == [('events', ['Event'])]
 
     def test_buffer_fills_up(self):
-        with patch('six.moves.urllib.request.urlopen', return_value=self.mock) as urlopen:
-            for i in range(self.MAX_LENGTH - 1):
-                self.consumer.send('events', '"Event"')
-                assert not self.mock.called
-
-            self.consumer.send('events', '"Last Event"')
+        for i in range(self.MAX_LENGTH - 1):
+            self.consumer.send('events', '"Event"')
+        assert len(self.log) == 0
 
-            assert urlopen.call_count == 1
-            ((request,), _) = urlopen.call_args
-            assert request.get_full_url() == 'https://api.mixpanel.com/track'
-            assert qs(request.data) == \
-                qs('ip=0&data=WyJFdmVudCIsIkV2ZW50IiwiRXZlbnQiLCJFdmVudCIsIkV2ZW50IiwiRXZlbnQiLCJFdmVudCIsIkV2ZW50IiwiRXZlbnQiLCJMYXN0IEV2ZW50Il0%3D&verbose=1')
+        self.consumer.send('events', '"Last Event"')
+        assert len(self.log) == 1
+        assert self.log == [('events', [
+            'Event', 'Event', 'Event', 'Event', 'Event',
+            'Event', 'Event', 'Event', 'Event', 'Last Event',
+        ])]
 
-    def test_unknown_endpoint(self):
+    def test_unknown_endpoint_raises_on_send(self):
+        # Ensure the exception isn't hidden until a flush.
         with pytest.raises(mixpanel.MixpanelException):
             self.consumer.send('unknown', '1')
 
@@ -394,17 +386,19 @@ class TestBufferedConsumer:
         error_mock = Mock()
         error_mock.read.return_value = six.b('{"status": 0, "error": "arbitrary error"}')
         broken_json = '{broken JSON'
+        consumer = mixpanel.BufferedConsumer(2)
         with patch('six.moves.urllib.request.urlopen', return_value=error_mock):
-            self.consumer.send('events', broken_json)
+            consumer.send('events', broken_json)
             with pytest.raises(mixpanel.MixpanelException) as excinfo:
-                self.consumer.flush()
+                consumer.flush()
             assert excinfo.value.message == '[%s]' % broken_json
             assert excinfo.value.endpoint == 'events'
 
-    def test_import_data_receives_api_key(self):
-        # Ensure BufferedConsumer.send accepts the API_KEY parameter needed for
-        # import_data; see #62.
+    def test_send_remembers_api_key(self):
         self.consumer.send('imports', '"Event"', api_key='MY_API_KEY')
+        assert len(self.log) == 0
+        self.consumer.flush()
+        assert self.log == [('imports', ['Event'], 'MY_API_KEY')]
 
 
 class TestFunctional: