Merge pull request #424 from eribeiro/411-connectstring-accept-list

KazooClient hosts list should accept a list of multiple endpoints.

2 files changed, 66 insertions, 4 deletions

diff --git a/kazoo/hosts.py b/kazoo/hosts.py
index 7f5aad5..04643dc 100644
--- a/kazoo/hosts.py
+++ b/kazoo/hosts.py
@@ -2,12 +2,22 @@ from six.moves import urllib_parse
 
 
 def collect_hosts(hosts):
-    """Collect a set of hosts and an optional chroot from a string."""
-    host_ports, chroot = hosts.partition("/")[::2]
-    chroot = "/" + chroot if chroot else None
+    """
+       Collect a set of hosts and an optional chroot from
+       a string or a list of strings.
+    """
+    if isinstance(hosts, list):
+        if hosts[-1].strip().startswith('/'):
+            host_ports, chroot = hosts[:-1], hosts[-1]
+        else:
+            host_ports, chroot = hosts, None
+    else:
+        host_ports, chroot = hosts.partition("/")[::2]
+        host_ports = host_ports.split(",")
+        chroot = "/" + chroot if chroot else None
 
     result = []
-    for host_port in host_ports.split(","):
+    for host_port in host_ports:
         # put all complexity of dealing with
         # IPv4 & IPv6 address:port on the urlsplit
         res = urllib_parse.urlsplit("xxx://" + host_port)
diff --git a/kazoo/tests/test_hosts.py b/kazoo/tests/test_hosts.py
new file mode 100644
index 0000000..cce623c
--- /dev/null
+++ b/kazoo/tests/test_hosts.py
@@ -0,0 +1,52 @@
+from unittest import TestCase
+from kazoo.hosts import collect_hosts
+
+
+class HostsTestCase(TestCase):
+
+    def test_ipv4(self):
+        hosts, chroot = collect_hosts('127.0.0.1:2181, 192.168.1.2:2181, \
+                                       132.254.111.10:2181')
+        self.assertEquals([('127.0.0.1', 2181),
+                           ('192.168.1.2', 2181),
+                           ('132.254.111.10', 2181)], hosts)
+        self.assertEquals(None, chroot)
+
+        hosts, chroot = collect_hosts(['127.0.0.1:2181',
+                                       '192.168.1.2:2181',
+                                       '132.254.111.10:2181'])
+        self.assertEquals([('127.0.0.1', 2181),
+                           ('192.168.1.2', 2181),
+                           ('132.254.111.10', 2181)], hosts)
+        self.assertEquals(None, chroot)
+
+    def test_ipv6(self):
+        hosts, chroot = collect_hosts('[fe80::200:5aee:feaa:20a2]:2181')
+        self.assertEquals([('fe80::200:5aee:feaa:20a2', 2181)], hosts)
+        self.assertEquals(None, chroot)
+
+        hosts, chroot = collect_hosts(['[fe80::200:5aee:feaa:20a2]:2181'])
+        self.assertEquals([('fe80::200:5aee:feaa:20a2', 2181)], hosts)
+        self.assertEquals(None, chroot)
+
+    def test_hosts_list(self):
+
+        hosts, chroot = collect_hosts('zk01:2181, zk02:2181, zk03:2181')
+        expected1 = [('zk01', 2181), ('zk02', 2181), ('zk03', 2181)]
+        self.assertEquals(expected1, hosts)
+        self.assertEquals(None, chroot)
+
+        hosts, chroot = collect_hosts(['zk01:2181', 'zk02:2181', 'zk03:2181'])
+        self.assertEquals(expected1, hosts)
+        self.assertEquals(None, chroot)
+
+        expected2 = '/test'
+        hosts, chroot = collect_hosts('zk01:2181, zk02:2181, zk03:2181/test')
+        self.assertEquals(expected1, hosts)
+        self.assertEquals(expected2, chroot)
+
+        hosts, chroot = collect_hosts(['zk01:2181',
+                                       'zk02:2181',
+                                       'zk03:2181', '/test'])
+        self.assertEquals(expected1, hosts)
+        self.assertEquals(expected2, chroot)