DataQueue message saving and flushing

txstatsd/protocol.py

@@ -48,8 +48,18 @@ class DataQueue(object):
     """Manages the queue of sent data, so that it can be really sent later when
     the host is resolved."""
 
+    def __init__(self):
+        self._queue = []
+
     def write(self, data, callback):
         """Queue the given data, so that it's sent later."""
+        self._queue.append((data, callback))
+
+    def flush(self):
+        """Flush the queue, returning its items."""
+        items = self._queue
+        self._queue = []
+        return items
 
 
 class TransportGateway(object):

txstatsd/tests/test_client.py

@@ -302,6 +302,35 @@ def test_sends_messages_to_queue_before_host_resolves(self):
             self.assertEqual(self.client.write(message, callback), None)
 
 
+class DataQueueTest(TestCase):
+    """Tests for the DataQueue class."""
+
+    def setUp(self):
+        super(DataQueueTest, self).setUp()
+        self.queue = DataQueue()
+
+    def test_queues_messages_and_callbacks(self):
+        """All messages are queued with their respective callbacks."""
+        self.queue.write(data=1, callback='1')
+        self.queue.write(data=2, callback='2')
+        self.queue.write(data=3, callback='3')
+
+        self.assertEqual(self.queue.flush(), [
+            (1, '1'),
+            (2, '2'),
+            (3, '3'),
+        ])
+
+    def test_flushes_the_queue(self):
+        """All messages are queued with their respective callbacks."""
+        self.queue.write(data=1, callback='1')
+        self.queue.write(data=2, callback='2')
+        self.queue.write(data=3, callback='3')
+
+        self.queue.flush()
+        self.assertEqual(self.queue.flush(), [])
+
+
 class TestConsistentHashingClient(TestCase):
 
     def test_hash_with_single_client(self):