More unit tests

Run like:

    python -m unittest -v test

1 files changed, 64 insertions, 2 deletions

diff --git a/test.py b/test.py
index 3f5908e..56f1e1e 100644
--- a/test.py
+++ b/test.py
@@ -1,9 +1,51 @@
-import binascii
+import os
+import random
+import struct
 import unittest
 
 from kafka import KafkaClient
+from kafka import gzip_encode, gzip_decode, length_prefix_message
+
+ITERATIONS = 1000
+STRLEN = 100
+
+def random_string():
+    return os.urandom(random.randint(0, STRLEN))
+
+class TestMisc(unittest.TestCase):
+    def test_length_prefix(self):
+        for i in xrange(ITERATIONS):
+            s1 = random_string()
+            s2 = length_prefix_message(s1)
+            self.assertEquals(struct.unpack('>i', s2[0:4])[0], len(s1)) 
+        
+class TestCodec(unittest.TestCase):
+    def test_gzip(self):
+        for i in xrange(ITERATIONS):
+            s1 = random_string()
+            s2 = gzip_decode(gzip_encode(s1))
+            self.assertEquals(s1,s2)
 
 class TestMessage(unittest.TestCase):
+    def test_create(self):
+        msg = KafkaClient.create_message("testing")
+        self.assertEquals(msg.payload, "testing")
+        self.assertEquals(msg.magic, 1)
+        self.assertEquals(msg.attributes, 0)
+        self.assertEquals(msg.crc, -386704890) 
+
+    def test_create_gzip(self):
+        msg = KafkaClient.create_gzip_message("testing")
+        self.assertEquals(msg.magic, 1)
+        self.assertEquals(msg.attributes, 1)
+        # Can't check the crc or payload for gzip since it's non-deterministic
+        (messages, _) = KafkaClient.read_message_set(gzip_decode(msg.payload))
+        inner = messages[0]
+        self.assertEquals(inner.magic, 1)
+        self.assertEquals(inner.attributes, 0)
+        self.assertEquals(inner.payload, "testing")
+        self.assertEquals(inner.crc, -386704890) 
+
     def test_message_simple(self):
         msg = KafkaClient.create_message("testing")
         enc = KafkaClient.encode_message(msg)
@@ -28,7 +70,6 @@ class TestMessage(unittest.TestCase):
         self.assertEquals(messages[0].payload, "one")
         self.assertEquals(messages[1].payload, "two")
         self.assertEquals(messages[2].payload, "three")
-         
 
     def test_message_gzip(self):
         msg = KafkaClient.create_gzip_message("one", "two", "three")
@@ -40,5 +81,26 @@ class TestMessage(unittest.TestCase):
         self.assertEquals(messages[1].payload, "two")
         self.assertEquals(messages[2].payload, "three")
 
+    def test_message_simple_random(self):
+        for i in xrange(ITERATIONS):
+            n = random.randint(0, 10)
+            msgs = [KafkaClient.create_message(random_string()) for j in range(n)]
+            enc = KafkaClient.encode_message_set(msgs)
+            (messages, read) = KafkaClient.read_message_set(enc)
+            self.assertEquals(len(messages), n)
+            for j in range(n):
+                self.assertEquals(messages[j], msgs[j])
+
+    def test_message_gzip_random(self):
+        for i in xrange(ITERATIONS):
+            n = random.randint(0, 10)
+            strings = [random_string() for j in range(n)]
+            msg = KafkaClient.create_gzip_message(*strings)
+            enc = KafkaClient.encode_message(msg)
+            (messages, read) = KafkaClient.read_message_set(enc)
+            self.assertEquals(len(messages), n)
+            for j in range(n):
+                self.assertEquals(messages[j].payload, strings[j])
+
 if __name__ == '__main__':
     unittest.main()