Finish breaking out integration tests

1 files changed, 16 insertions, 2 deletions

diff --git a/test/testutil.py b/test/testutil.py
index 4866b9d..2cf62eb 100644
--- a/test/testutil.py
+++ b/test/testutil.py
@@ -26,6 +26,7 @@ def ensure_topic_creation(client, topic_name, timeout = 30):
         time.sleep(1)
 
 class KafkaIntegrationTestCase(unittest.TestCase):
+    create_client = True
     topic = None
 
     def setUp(self):
@@ -33,13 +34,17 @@ class KafkaIntegrationTestCase(unittest.TestCase):
         if not self.topic:
             self.topic = "%s-%s" % (self.id()[self.id().rindex(".") + 1:], random_string(10))
 
-        self.client = KafkaClient('%s:%d' % (self.server.host, self.server.port))
+        if self.create_client:
+            self.client = KafkaClient('%s:%d' % (self.server.host, self.server.port))
+
         ensure_topic_creation(self.client, self.topic)
+
         self._messages = {}
 
     def tearDown(self):
         super(KafkaIntegrationTestCase, self).tearDown()
-        self.client.close()
+        if self.create_client:
+            self.client.close()
 
     def current_offset(self, topic, partition):
         offsets, = self.client.send_offset_request([ OffsetRequest(topic, partition, -1, 1) ])
@@ -54,4 +59,13 @@ class KafkaIntegrationTestCase(unittest.TestCase):
 
         return self._messages[s]
 
+class Timer(object):
+    def __enter__(self):
+        self.start = time.time()
+        return self
+
+    def __exit__(self, *args):
+        self.end = time.time()
+        self.interval = self.end - self.start
+
 logging.basicConfig(level=logging.DEBUG)