Rearrange connection tests to separate legacy KafkaConnection

1 files changed, 57 insertions, 9 deletions

diff --git a/test/test_conn.py b/test/test_conn.py
index 6a3b154..4f2b12f 100644
--- a/test/test_conn.py
+++ b/test/test_conn.py
@@ -7,7 +7,7 @@ import time
 
 import pytest
 
-from kafka.conn import BrokerConnection, ConnectionStates
+from kafka.conn import BrokerConnection, ConnectionStates, collect_hosts
 from kafka.protocol.api import RequestHeader
 from kafka.protocol.metadata import MetadataRequest
 
@@ -29,14 +29,14 @@ def conn(_socket):
 
 
 @pytest.mark.parametrize("states", [
-  (([EINPROGRESS, EALREADY], ConnectionStates.CONNECTING),),
-  (([EALREADY, EALREADY], ConnectionStates.CONNECTING),),
-  (([0], ConnectionStates.CONNECTED),),
-  (([EINPROGRESS, EALREADY], ConnectionStates.CONNECTING),
-   ([ECONNRESET], ConnectionStates.DISCONNECTED)),
-  (([EINPROGRESS, EALREADY], ConnectionStates.CONNECTING),
-   ([EALREADY], ConnectionStates.CONNECTING),
-   ([EISCONN], ConnectionStates.CONNECTED)),
+    (([EINPROGRESS, EALREADY], ConnectionStates.CONNECTING),),
+    (([EALREADY, EALREADY], ConnectionStates.CONNECTING),),
+    (([0], ConnectionStates.CONNECTED),),
+    (([EINPROGRESS, EALREADY], ConnectionStates.CONNECTING),
+     ([ECONNRESET], ConnectionStates.DISCONNECTED)),
+    (([EINPROGRESS, EALREADY], ConnectionStates.CONNECTING),
+     ([EALREADY], ConnectionStates.CONNECTING),
+     ([EISCONN], ConnectionStates.CONNECTED)),
 ])
 def test_connect(_socket, conn, states):
     assert conn.state is ConnectionStates.DISCONNECTED
@@ -216,3 +216,51 @@ def test_recv(_socket, conn):
 
 def test_close(conn):
     pass # TODO
+
+
+def test_collect_hosts__happy_path():
+    hosts = "127.0.0.1:1234,127.0.0.1"
+    results = collect_hosts(hosts)
+    assert set(results) == set([
+        ('127.0.0.1', 1234, socket.AF_INET),
+        ('127.0.0.1', 9092, socket.AF_INET),
+    ])
+
+
+def test_collect_hosts__ipv6():
+    hosts = "[localhost]:1234,[2001:1000:2000::1],[2001:1000:2000::1]:1234"
+    results = collect_hosts(hosts)
+    assert 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():
+    hosts = [
+        'localhost:1234',
+        'localhost',
+        '[localhost]',
+        '2001::1',
+        '[2001::1]',
+        '[2001::1]:1234',
+    ]
+    results = collect_hosts(hosts)
+    assert set(results) == set([
+        ('localhost', 1234, socket.AF_UNSPEC),
+        ('localhost', 9092, socket.AF_UNSPEC),
+        ('localhost', 9092, socket.AF_INET6),
+        ('2001::1', 9092, socket.AF_INET6),
+        ('2001::1', 9092, socket.AF_INET6),
+        ('2001::1', 1234, socket.AF_INET6),
+    ])
+
+
+def test_collect_hosts__with_spaces():
+    hosts = "localhost:1234, localhost"
+    results = collect_hosts(hosts)
+    assert set(results) == set([
+        ('localhost', 1234, socket.AF_UNSPEC),
+        ('localhost', 9092, socket.AF_UNSPEC),
+    ])