Add tests (#3)

* Add tests

* Move buffer into its own module

* Move test

* Pep8

* Pep8

Author: mdelaurentis

stitchclient/buffer.py

@@ -0,0 +1,62 @@
+import time
+import collections
+
+from collections import deque
+
+BufferEntry = collections.namedtuple(
+    'BufferEntry',
+    'timestamp value callback_arg')
+
+MAX_BATCH_SIZE_BYTES = 4194304
+MAX_MESSAGES_PER_BATCH = 10000
+
+
+class Buffer(object):
+
+    def __init__(self):
+        self._queue = deque()
+        self._available_bytes = 0
+
+    def put(self, value, callback_arg):
+        # We need two extra bytes for the [ and ] wrapping the record.
+        max_len = MAX_BATCH_SIZE_BYTES - 2
+
+        if len(value) > max_len:
+            raise ValueError(
+                "Can't accept a record larger than {} bytes".format(max_len))
+
+        self._queue.append(BufferEntry(timestamp=time.time()*1000,
+                                       value=value,
+                                       callback_arg=callback_arg))
+        self._available_bytes += len(value.encode("utf8"))
+
+    def take(self, batch_size_bytes, batch_delay_millis):
+        if len(self._queue) == 0:
+            return None
+
+        t = time.time() * 1000
+        t0 = self._queue[0].timestamp
+        enough_bytes = self._available_bytes >= batch_size_bytes
+        enough_messages = len(self._queue) >= MAX_MESSAGES_PER_BATCH
+        enough_time = t - t0 >= batch_delay_millis
+        ready = enough_bytes or enough_messages or enough_time
+
+        if not ready:
+            return None
+
+        entries = []
+        size = 2
+
+        while (len(self._queue) > 0 and
+               size + len(self._queue[0].value.encode("utf8")) <
+               MAX_BATCH_SIZE_BYTES):
+            entry = self._queue.popleft()
+
+            # add one for the comma that will be needed to link entries
+            # together
+            entry_size = len(entry.value.encode("utf8"))
+            size += entry_size + 1
+            self._available_bytes -= entry_size
+            entries.append(entry)
+
+        return entries

stitchclient/client.py

@@ -1,11 +1,8 @@
 import os
-import time
 import logging
-import collections
-
 import requests
+from .buffer import Buffer
 
-from collections import deque
 from io import StringIO
 from transit.writer import Writer
 from transit.reader import Reader
@@ -16,48 +13,9 @@
 DEFAULT_BATCH_DELAY_MILLIS = 60000
 DEFAULT_STITCH_URL = 'https://api.stitchdata.com/v2/import/push'
 
-BufferEntry = collections.namedtuple('BufferEntry', 'timestamp value callback_arg')
 
 class Client(object):
 
-    class Buffer(object):
-
-        MAX_BATCH_SIZE_BYTES = 4194304
-        MAX_MESSAGES_PER_BATCH = 10000
-
-        _queue = deque()
-        _available_bytes = 0
-
-        def put(self, value, callback_arg):
-            self._queue.append(BufferEntry(timestamp=time.time()*1000, value=value, callback_arg=callback_arg))
-            self._available_bytes += len(value.encode("utf8"))
-
-        def take(self, batch_size_bytes, batch_delay_millis):
-            if len(self._queue) == 0:
-                return None
-
-            ready = self._available_bytes >= batch_size_bytes or \
-                    len(self._queue) >= self.MAX_MESSAGES_PER_BATCH or \
-                    time.time()*1000 - self._queue[0].timestamp >= batch_delay_millis
-
-            if not ready:
-                return None
-
-            entries = []
-            size = 2
-
-            while len(self._queue) > 0 and \
-                  size + len(self._queue[0].value.encode("utf8")) < self.MAX_BATCH_SIZE_BYTES:
-                entry = self._queue.popleft()
-
-                # add one for the comma that will be needed to link entries together
-                entry_size = len(entry.value.encode("utf8"))
-                size += entry_size + 1
-                self._available_bytes -= entry_size
-                entries.append(entry)
-
-            return entries
-
     _buffer = Buffer()
 
     def __init__(self,
@@ -70,7 +28,8 @@ def __init__(self,
                  batch_size_bytes=DEFAULT_BATCH_SIZE_BYTES,
                  batch_delay_millis=DEFAULT_BATCH_DELAY_MILLIS):
 
-        assert isinstance(client_id, int), 'client_id is not an integer: {}'.format(client_id)
+        assert(isinstance(client_id, int),
+               'client_id is not an integer: {}'.format(client_id))
 
         self.client_id = client_id
         self.token = token
@@ -81,7 +40,7 @@ def __init__(self,
         self.batch_delay_millis = batch_delay_millis
         self.callback_function = callback_function
 
-    def push(self, message, callback_arg = None):
+    def push(self, message, callback_arg=None):
         """
         message must be a dict with at least these keys:
             action, table_name, key_names, sequence, data
@@ -91,10 +50,8 @@ def push(self, message, callback_arg = None):
 
         if message['action'] == 'upsert':
             message.setdefault('key_names', self.key_names)
-        elif message['action'] == 'switch_view':
-            pass
         else:
-            raise ValueError('Message action property must be one of: "upsert", "switch_view"')
+            raise ValueError('Message action property must be "upsert"')
 
         message['client_id'] = self.client_id
         message.setdefault('table_name', self.table_name)
@@ -104,7 +61,8 @@ def push(self, message, callback_arg = None):
             writer.write(message)
             self._buffer.put(s.getvalue(), callback_arg)
 
-        batch = self._buffer.take(self.batch_size_bytes, self.batch_delay_millis)
+        batch = self._buffer.take(
+            self.batch_size_bytes, self.batch_delay_millis)
         if batch is not None:
             self._send_batch(batch)
 
@@ -133,12 +91,12 @@ def _send_batch(self, batch):
             if self.callback_function is not None:
                 self.callback_function([x.callback_arg for x in batch])
         else:
-            raise RuntimeError("Error sending data to the Stitch API. {0.status_code} - {0.content}"
+            raise RuntimeError("Error sending data to the Stitch API. {0.status_code} - {0.content}"  # nopep8
                                .format(response))
 
     def flush(self):
         while True:
-            batch = self._buffer.take(0,0)
+            batch = self._buffer.take(0, 0)
             if batch is None:
                 return
 
@@ -153,8 +111,10 @@ def __exit__(self, exception_type, exception_value, traceback):
 if __name__ == "__main__":
     logging.basicConfig(level=logging.DEBUG)
 
-    with Client(int(os.environ['STITCH_CLIENT_ID']), os.environ['STITCH_TOKEN'], callback_function=print) as c:
-        for i in range(1,10):
+    with Client(int(os.environ['STITCH_CLIENT_ID']),
+                os.environ['STITCH_TOKEN'],
+                callback_function=print) as c:
+        for i in range(1, 10):
             c.push({'action': 'upsert',
                     'table_name': 'test_table',
                     'key_names': ['id'],

tests/__init__.py

@@ -0,0 +1,101 @@
+import unittest
+import stitchclient.client
+import time
+from stitchclient.buffer import MAX_BATCH_SIZE_BYTES
+
+class TestTargetStitch(unittest.TestCase):
+
+    client_id = 1
+    token = 'asdf'
+    table_name = 'users'
+    key_names = ['id']
+
+    def test_push(self):
+        pass
+        # client = stitchclient.Client(
+        #     client_id=client_id,
+        #     token=token,
+        #     table_name=table_name,
+        #     key_names=key_names,
+        #     batch_size_bytes=1000,
+        #     batch_delay_millis=100000)
+
+
+tiny_record = 'apple'
+big_record  = 'a' * 1500000
+huge_record = 'a' * 5000000
+
+class TestBuffer(unittest.TestCase):
+
+    buffer = None
+
+        
+    def test_single_record_available_immediately(self):
+        buf = stitchclient.client.Buffer()
+        buf.put(tiny_record, None)
+        self.assertEqual(buf.take(0, 0)[0].value,
+                         tiny_record)
+
+    def test_withhold_until_bytes_available(self):
+        buf = stitchclient.client.Buffer()        
+        batch_size_bytes = int(len(tiny_record) * 2 + len(tiny_record) / 2.0)
+        batch_delay_millis = 1000000000
+        put = lambda: buf.put(tiny_record, None)
+        take = lambda: buf.take(batch_size_bytes, batch_delay_millis)
+        
+        put()
+        self.assertTrue(take() is None)
+        put()
+        self.assertTrue(take() is None)
+        put()
+        res = take()
+        self.assertEqual([x.value for x in res], ['apple', 'apple', 'apple'])
+        
+    def test_buffer_empty_after_batch(self):
+        buf = stitchclient.client.Buffer()        
+        put = lambda: buf.put(tiny_record, None)
+        take = lambda: buf.take(0, 0)        
+        put()
+        put()
+        put()
+        self.assertIsNotNone(take())
+        self.assertIsNone(take())
+
+    def test_does_not_exceed_max_batch_size(self):
+        buf = stitchclient.client.Buffer()        
+        put = lambda: buf.put(big_record, None)
+        take = lambda: buf.take(0, 0)
+
+        put()
+        put()
+        put()
+
+        b1 = take()
+        b2 = take()
+        b3 = take()
+
+        b1len = 0
+        b2len = 0
+
+        for x in b1:
+            b1len += len(x.value)
+        for x in b2:
+            b2len += len(x.value)            
+        self.assertTrue(b1len < MAX_BATCH_SIZE_BYTES)
+        self.assertTrue(b2len < b1len)
+        self.assertIsNone(b3)
+
+    def test_cant_put_record_larger_than_max_message_size(self):
+        buf = stitchclient.client.Buffer()        
+        with self.assertRaises(ValueError):
+            buf.put(huge_record, None)
+
+    def test_trigger_batch_at_10k_messages(self):
+        buf = stitchclient.client.Buffer()        
+        put = lambda: buf.put(tiny_record, None)
+        take = lambda: buf.take(MAX_BATCH_SIZE_BYTES, 60000)
+        for i in range(9999):
+            put()
+        self.assertTrue(take() is None)
+        put()
+        self.assertIsNotNone(take())