Merge pull request #615 from TimEvens/master

Kafka IPv6 Support.

3 files changed, 37 insertions, 18 deletions

diff --git a/test/test_client.py b/test/test_client.py
index a53fce1..6980434 100644
--- a/test/test_client.py
+++ b/test/test_client.py
@@ -41,7 +41,7 @@ class TestSimpleClient(unittest.TestCase):
             client = SimpleClient(hosts=['kafka01:9092', 'kafka02:9092', 'kafka03:9092'])
 
         self.assertEqual(
-            sorted([('kafka01', 9092), ('kafka02', 9092), ('kafka03', 9092)]),
+            sorted([('kafka01', 9092, socket.AF_INET), ('kafka02', 9092, socket.AF_INET), ('kafka03', 9092, socket.AF_INET)]),
             sorted(client.hosts))
 
     def test_init_with_csv(self):
@@ -49,7 +49,7 @@ class TestSimpleClient(unittest.TestCase):
             client = SimpleClient(hosts='kafka01:9092,kafka02:9092,kafka03:9092')
 
         self.assertEqual(
-            sorted([('kafka01', 9092), ('kafka02', 9092), ('kafka03', 9092)]),
+            sorted([('kafka01', 9092, socket.AF_INET), ('kafka02', 9092, socket.AF_INET), ('kafka03', 9092, socket.AF_INET)]),
             sorted(client.hosts))
 
     def test_init_with_unicode_csv(self):
@@ -57,7 +57,7 @@ class TestSimpleClient(unittest.TestCase):
             client = SimpleClient(hosts=u'kafka01:9092,kafka02:9092,kafka03:9092')
 
         self.assertEqual(
-            sorted([('kafka01', 9092), ('kafka02', 9092), ('kafka03', 9092)]),
+            sorted([('kafka01', 9092, socket.AF_INET), ('kafka02', 9092, socket.AF_INET), ('kafka03', 9092, socket.AF_INET)]),
             sorted(client.hosts))
 
     @patch.object(SimpleClient, '_get_conn')
@@ -70,7 +70,7 @@ class TestSimpleClient(unittest.TestCase):
         for val in mocked_conns.values():
             mock_conn(val, success=False)
 
-        def mock_get_conn(host, port):
+        def mock_get_conn(host, port, afi):
             return mocked_conns[(host, port)]
         conn.side_effect = mock_get_conn
 
@@ -98,7 +98,7 @@ class TestSimpleClient(unittest.TestCase):
         mocked_conns[('kafka02', 9092)].send.return_value = future
         mocked_conns[('kafka02', 9092)].recv.side_effect = lambda: future.success('valid response')
 
-        def mock_get_conn(host, port):
+        def mock_get_conn(host, port, afi):
             return mocked_conns[(host, port)]
 
         # patch to avoid making requests before we want it
@@ -409,3 +409,4 @@ class TestSimpleClient(unittest.TestCase):
             self.assertEqual(big_num + 1, client._next_id())
             self.assertEqual(big_num + 2, client._next_id())
             self.assertEqual(0, client._next_id())
+
diff --git a/test/test_client_async.py b/test/test_client_async.py
index 2e0d9b4..e0b98c4 100644
--- a/test/test_client_async.py
+++ b/test/test_client_async.py
@@ -1,4 +1,5 @@
 import time
+import socket
 
 import pytest
 
@@ -12,11 +13,11 @@ from kafka.protocol.produce import ProduceRequest
 
 
 @pytest.mark.parametrize("bootstrap,expected_hosts", [
-    (None, [('localhost', 9092)]),
-    ('foobar:1234', [('foobar', 1234)]),
-    ('fizzbuzz', [('fizzbuzz', 9092)]),
-    ('foo:12,bar:34', [('foo', 12), ('bar', 34)]),
-    (['fizz:56', 'buzz'], [('fizz', 56), ('buzz', 9092)]),
+    (None, [('localhost', 9092, socket.AF_INET)]),
+    ('foobar:1234', [('foobar', 1234, socket.AF_INET)]),
+    ('fizzbuzz', [('fizzbuzz', 9092, socket.AF_INET)]),
+    ('foo:12,bar:34', [('foo', 12, socket.AF_INET), ('bar', 34, socket.AF_INET)]),
+    (['fizz:56', 'buzz'], [('fizz', 56, socket.AF_INET), ('buzz', 9092, socket.AF_INET)]),
 ])
 def test_bootstrap_servers(mocker, bootstrap, expected_hosts):
     mocker.patch.object(KafkaClient, '_bootstrap')
@@ -47,7 +48,7 @@ def conn(mocker):
 def test_bootstrap_success(conn):
     conn.state = ConnectionStates.CONNECTED
     cli = KafkaClient()
-    conn.assert_called_once_with('localhost', 9092, **cli.config)
+    conn.assert_called_once_with('localhost', 9092, socket.AF_INET, **cli.config)
     conn.connect.assert_called_with()
     conn.send.assert_called_once_with(MetadataRequest([]))
     assert cli._bootstrap_fails == 0
@@ -57,7 +58,7 @@ def test_bootstrap_success(conn):
 def test_bootstrap_failure(conn):
     conn.state = ConnectionStates.DISCONNECTED
     cli = KafkaClient()
-    conn.assert_called_once_with('localhost', 9092, **cli.config)
+    conn.assert_called_once_with('localhost', 9092, socket.AF_INET, **cli.config)
     conn.connect.assert_called_with()
     conn.close.assert_called_with()
     assert cli._bootstrap_fails == 1
diff --git a/test/test_conn.py b/test/test_conn.py
index 684ffe5..f0ef8fb 100644
--- a/test/test_conn.py
+++ b/test/test_conn.py
@@ -51,21 +51,37 @@ class ConnTest(unittest.TestCase):
         results = collect_hosts(hosts)
 
         self.assertEqual(set(results), set([
-            ('localhost', 1234),
-            ('localhost', 9092),
+            ('localhost', 1234, socket.AF_INET),
+            ('localhost', 9092, socket.AF_INET),
+        ]))
+
+    def test_collect_hosts__ipv6(self):
+        hosts = "[localhost]:1234,[2001:1000:2000::1],[2001:1000:2000::1]:1234"
+        results = collect_hosts(hosts)
+
+        self.assertEqual(set(results), set([
+            ('localhost', 1234, socket.AF_INET6),
+            ('2001:1000:2000::1', 9092, socket.AF_INET6),
+            ('2001:1000:2000::1', 1234, socket.AF_INET6),
         ]))
 
     def test_collect_hosts__string_list(self):
         hosts = [
             'localhost:1234',
             'localhost',
+            '[localhost]',
+            '2001::1',
+            '[2001::1]:1234',
         ]
 
         results = collect_hosts(hosts)
 
         self.assertEqual(set(results), set([
-            ('localhost', 1234),
-            ('localhost', 9092),
+            ('localhost', 1234, socket.AF_INET),
+            ('localhost', 9092, socket.AF_INET),
+            ('localhost', 9092, socket.AF_INET6),
+            ('2001::1', 9092, socket.AF_INET6),
+            ('2001::1', 1234, socket.AF_INET6),
         ]))
 
     def test_collect_hosts__with_spaces(self):
@@ -73,10 +89,11 @@ class ConnTest(unittest.TestCase):
         results = collect_hosts(hosts)
 
         self.assertEqual(set(results), set([
-            ('localhost', 1234),
-            ('localhost', 9092),
+            ('localhost', 1234, socket.AF_INET),
+            ('localhost', 9092, socket.AF_INET),
         ]))
 
+
     def test_send(self):
         self.conn.send(self.config['request_id'], self.config['payload'])
         self.conn._sock.sendall.assert_called_with(self.config['payload'])