Merge "Allow testing with external zookeeper"0.3.0

3 files changed, 55 insertions, 4 deletions

diff --git a/nodepool/tests/__init__.py b/nodepool/tests/__init__.py
index ab4b7a9..afae79e 100644
--- a/nodepool/tests/__init__.py
+++ b/nodepool/tests/__init__.py
@@ -29,6 +29,7 @@ import time
 
 import fixtures
 import gear
+import kazoo.client
 import testresources
 import testtools
 
@@ -116,6 +117,24 @@ class ZookeeperServerFixture(fixtures.Fixture):
     def setUp(self):
         super(ZookeeperServerFixture, self).setUp()
 
+        if 'NODEPOOL_ZK_HOST' in os.environ:
+           if ':' in os.environ['NODEPOOL_ZK_HOST']:
+               host, port = os.environ['NODEPOOL_ZK_HOST'].split(':')
+           else:
+               host = os.environ['NODEPOOL_ZK_HOST']
+               port = None
+
+           self.zookeeper_host = host
+
+           if not port:
+               self.zookeeper_port = 2181
+           else:
+               self.zookeeper_port = int(port)
+
+           return
+
+        self.zookeeper_host = '127.0.0.1'
+
         # Get the local port range, we're going to pick one at a time
         # at random to try.
         with open('/proc/sys/net/ipv4/ip_local_port_range') as f:
@@ -500,4 +519,38 @@ class ZKTestCase(BaseTestCase):
         super(ZKTestCase, self).setUp()
         f = ZookeeperServerFixture()
         self.useFixture(f)
+        self.zookeeper_host = f.zookeeper_host
         self.zookeeper_port = f.zookeeper_port
+        self.chroot_path = "/nodepool_test/%s" % self.getUniqueInteger()
+
+        # Ensure the chroot path exists and clean up an pre-existing znodes
+        _tmp_client = kazoo.client.KazooClient(
+            hosts='%s:%s' % (self.zookeeper_host, self.zookeeper_port))
+        _tmp_client.start()
+
+        if _tmp_client.exists(self.chroot_path):
+            _tmp_client.delete(self.chroot_path, recursive=True)
+
+        _tmp_client.ensure_path(self.chroot_path)
+        _tmp_client.stop()
+
+        # Create a chroot'ed client
+        self.zkclient = kazoo.client.KazooClient(
+            hosts='%s:%s%s' % (self.zookeeper_host,
+                               self.zookeeper_port,
+                               self.chroot_path)
+        )
+        self.zkclient.start()
+
+        self.addCleanup(self._cleanup)
+
+    def _cleanup(self):
+        '''Stop the client and remove the chroot path.'''
+        self.zkclient.stop()
+
+        # Need a non-chroot'ed client to remove the chroot path
+        _tmp_client = kazoo.client.KazooClient(
+            hosts='%s:%s' % (self.zookeeper_host, self.zookeeper_port))
+        _tmp_client.start()
+        _tmp_client.delete(self.chroot_path, recursive=True)
+        _tmp_client.stop()
diff --git a/nodepool/tests/test_builder.py b/nodepool/tests/test_builder.py
index 0cc1af7..622cc0c 100644
--- a/nodepool/tests/test_builder.py
+++ b/nodepool/tests/test_builder.py
@@ -18,7 +18,6 @@ import time
 import threading
 
 import fixtures
-import kazoo.client
 
 from nodepool import builder, exceptions, fakeprovider, tests
 
@@ -151,6 +150,4 @@ class TestNodepoolBuilder(tests.DBTestCase):
 
 class TestZookeeper(tests.ZKTestCase):
     def test_zk(self):
-        zk = kazoo.client.KazooClient(hosts='127.0.0.1:%s' % self.zookeeper_port)
-        zk.start()
-        zk.get('/')
+        self.zkclient.get('/')
diff --git a/tox.ini b/tox.ini
index 9052863..b2bafeb 100644
--- a/tox.ini
+++ b/tox.ini
@@ -12,6 +12,7 @@ usedevelop = True
 install_command = pip install {opts} {packages}
 deps = -r{toxinidir}/requirements.txt
        -r{toxinidir}/test-requirements.txt
+passenv = NODEPOOL_ZK_HOST
 commands =
   python setup.py testr --slowest --testr-args='{posargs}'