import json
import time
import unittest
import mock
try:
from queue import Queue
except ImportError:
from Queue import Queue
from posthog.consumer import MAX_MSG_SIZE, Consumer
from posthog.request import APIError
from posthog.test.test_utils import TEST_API_KEY
class TestConsumer(unittest.TestCase):
def test_next(self):
q = Queue()
consumer = Consumer(q, "")
q.put(1)
next = consumer.next()
self.assertEqual(next, [1])
def test_next_limit(self):
q = Queue()
flush_at = 50
consumer = Consumer(q, "", flush_at)
for i in range(10000):
q.put(i)
next = consumer.next()
self.assertEqual(next, list(range(flush_at)))
def test_dropping_oversize_msg(self):
q = Queue()
consumer = Consumer(q, "")
oversize_msg = {"m": "x" * MAX_MSG_SIZE}
q.put(oversize_msg)
next = consumer.next()
self.assertEqual(next, [])
self.assertTrue(q.empty())
def test_upload(self):
q = Queue()
consumer = Consumer(q, TEST_API_KEY)
track = {"type": "track", "event": "python event", "distinct_id": "distinct_id"}
q.put(track)
success = consumer.upload()
self.assertTrue(success)
def test_flush_interval(self):
# Put _n_ items in the queue, pausing a little bit more than
# _flush_interval_ after each one.
# The consumer should upload _n_ times.
q = Queue()
flush_interval = 0.3
consumer = Consumer(q, TEST_API_KEY, flush_at=10, flush_interval=flush_interval)
with mock.patch("posthog.consumer.batch_post") as mock_post:
consumer.start()
for i in range(0, 3):
track = {"type": "track", "event": "python event %d" % i, "distinct_id": "distinct_id"}
q.put(track)
time.sleep(flush_interval * 1.1)
self.assertEqual(mock_post.call_count, 3)
def test_multiple_uploads_per_interval(self):
# Put _flush_at*2_ items in the queue at once, then pause for
# _flush_interval_. The consumer should upload 2 times.
q = Queue()
flush_interval = 0.5
flush_at = 10
consumer = Consumer(q, TEST_API_KEY, flush_at=flush_at, flush_interval=flush_interval)
with mock.patch("posthog.consumer.batch_post") as mock_post:
consumer.start()
for i in range(0, flush_at * 2):
track = {"type": "track", "event": "python event %d" % i, "distinct_id": "distinct_id"}
q.put(track)
time.sleep(flush_interval * 1.1)
self.assertEqual(mock_post.call_count, 2)
def test_request(self):
consumer = Consumer(None, TEST_API_KEY)
track = {"type": "track", "event": "python event", "distinct_id": "distinct_id"}
consumer.request([track])
def _test_request_retry(self, consumer, expected_exception, exception_count):
def mock_post(*args, **kwargs):
mock_post.call_count += 1
if mock_post.call_count <= exception_count:
raise expected_exception
mock_post.call_count = 0
with mock.patch("posthog.consumer.batch_post", mock.Mock(side_effect=mock_post)):
track = {"type": "track", "event": "python event", "distinct_id": "distinct_id"}
# request() should succeed if the number of exceptions raised is
# less than the retries paramater.
if exception_count <= consumer.retries:
consumer.request([track])
else:
# if exceptions are raised more times than the retries
# parameter, we expect the exception to be returned to
# the caller.
try:
consumer.request([track])
except type(expected_exception) as exc:
self.assertEqual(exc, expected_exception)
else:
self.fail(
"request() should raise an exception if still failing after %d retries" % consumer.retries
)
def test_request_retry(self):
# we should retry on general errors
consumer = Consumer(None, TEST_API_KEY)
self._test_request_retry(consumer, Exception("generic exception"), 2)
# we should retry on server errors
consumer = Consumer(None, TEST_API_KEY)
self._test_request_retry(consumer, APIError(500, "Internal Server Error"), 2)
# we should retry on HTTP 429 errors
consumer = Consumer(None, TEST_API_KEY)
self._test_request_retry(consumer, APIError(429, "Too Many Requests"), 2)
# we should NOT retry on other client errors
consumer = Consumer(None, TEST_API_KEY)
api_error = APIError(400, "Client Errors")
try:
self._test_request_retry(consumer, api_error, 1)
except APIError:
pass
else:
self.fail("request() should not retry on client errors")
# test for number of exceptions raise > retries value
consumer = Consumer(None, TEST_API_KEY, retries=3)
self._test_request_retry(consumer, APIError(500, "Internal Server Error"), 3)
def test_pause(self):
consumer = Consumer(None, TEST_API_KEY)
consumer.pause()
self.assertFalse(consumer.running)
def test_max_batch_size(self):
q = Queue()
consumer = Consumer(q, TEST_API_KEY, flush_at=100000, flush_interval=3)
properties = {}
for n in range(0, 500):
properties[str(n)] = "one_long_property_value_to_build_a_big_event"
track = {"type": "track", "event": "python event", "distinct_id": "distinct_id", "properties": properties}
msg_size = len(json.dumps(track).encode())
# Let's capture 8MB of data to trigger two batches
n_msgs = int(8_000_000 / msg_size)
def mock_post_fn(_, data, **kwargs):
res = mock.Mock()
res.status_code = 200
request_size = len(data.encode())
# Batches close after the first message bringing it bigger than BATCH_SIZE_LIMIT, let's add 10% of margin
self.assertTrue(request_size < (5 * 1024 * 1024) * 1.1, "batch size (%d) higher than limit" % request_size)
return res
with mock.patch("posthog.request._session.post", side_effect=mock_post_fn) as mock_post:
consumer.start()
for _ in range(0, n_msgs + 2):
q.put(track)
q.join()
self.assertEqual(mock_post.call_count, 2)