Support setting kafka instance port explicitly in fixture

dpkp · dpkp · commit 0d1ccc7485ac · 2016-02-15T21:56:10.000-08:00
diff --git a/test/fixtures.py b/test/fixtures.py
@@ -151,16 +151,20 @@ def close(self):
 
 class KafkaFixture(Fixture):
     @classmethod
-    def instance(cls, broker_id, zk_host, zk_port, zk_chroot=None, replicas=1, partitions=2):
+    def instance(cls, broker_id, zk_host, zk_port,
+                 zk_chroot=None, port=None, replicas=1, partitions=2):
         if zk_chroot is None:
             zk_chroot = "kafka-python_" + str(uuid.uuid4()).replace("-", "_")
         if "KAFKA_URI" in os.environ:
             parse = urlparse(os.environ["KAFKA_URI"])
             (host, port) = (parse.hostname, parse.port)
             fixture = ExternalService(host, port)
         else:
-            (host, port) = ("127.0.0.1", get_open_port())
-            fixture = KafkaFixture(host, port, broker_id, zk_host, zk_port, zk_chroot, replicas, partitions)
+            if port is None:
+                port = get_open_port()
+            host = "127.0.0.1"
+            fixture = KafkaFixture(host, port, broker_id, zk_host, zk_port, zk_chroot,
+                                   replicas=replicas, partitions=partitions)
             fixture.open()
         return fixture
 
diff --git a/test/test_consumer_integration.py b/test/test_consumer_integration.py
@@ -27,8 +27,10 @@ def setUpClass(cls):
 
         cls.zk = ZookeeperFixture.instance()
         chroot = random_string(10)
-        cls.server1 = KafkaFixture.instance(0, cls.zk.host, cls.zk.port, chroot)
-        cls.server2 = KafkaFixture.instance(1, cls.zk.host, cls.zk.port, chroot)
+        cls.server1 = KafkaFixture.instance(0, cls.zk.host, cls.zk.port,
+                                            zk_chroot=chroot)
+        cls.server2 = KafkaFixture.instance(1, cls.zk.host, cls.zk.port,
+                                            zk_chroot=chroot)
 
         cls.server = cls.server1 # Bootstrapping server
 
diff --git a/test/test_failover_integration.py b/test/test_failover_integration.py
@@ -28,8 +28,11 @@ def setUp(self):
 
         # mini zookeeper, 3 kafka brokers
         self.zk = ZookeeperFixture.instance()
-        kk_args = [self.zk.host, self.zk.port, zk_chroot, replicas, partitions]
-        self.brokers = [KafkaFixture.instance(i, *kk_args) for i in range(replicas)]
+        kk_args = [self.zk.host, self.zk.port]
+        kk_kwargs = {'zk_chroot': zk_chroot, 'replicas': replicas,
+                     'partitions': partitions}
+        self.brokers = [KafkaFixture.instance(i, *kk_args, **kk_kwargs)
+                        for i in range(replicas)]
 
         hosts = ['%s:%d' % (b.host, b.port) for b in self.brokers]
         self.client = SimpleClient(hosts, timeout=2)
