From a0dca6ae128ea85213917e9141e1d4af52490711 Mon Sep 17 00:00:00 2001
From: David Shrewsbury <shrewsbury.dave@gmail.com>
Date: Thu, 7 Jul 2016 16:21:26 -0400
Subject: [PATCH] Allow testing with external zookeeper

It can be handy to test against an already running zookeeper, rather
than requiring one be installed in your test environment. Allow for
that by using the NODEPOOL_ZK_HOST environment variable.

This also gives each ZKTestCase based test a pre-existing client
connection that can be used for testing. To avoid znode path conflicts,
each connection is chroot'ed under /nodepool_test/{uniqueInt}. We
attempt to remove the {uniqueInt} znode and children as part of the
test cleanup (also removed at setUp time if cleanup fails). This will
leave a single znode for /nodepool_test, but if many tests fail for
some reason, it is easier to remove a single root node rather than
potentially many of them if we do not chroot.

Change-Id: I4fd61b16353c16bed2744e0d863ce188019430cd
---
 nodepool/tests/__init__.py     | 53 ++++++++++++++++++++++++++++++++++
 nodepool/tests/test_builder.py |  5 +---
 tox.ini                        |  1 +
 3 files changed, 55 insertions(+), 4 deletions(-)

diff --git a/nodepool/tests/__init__.py b/nodepool/tests/__init__.py
index ab4b7a9b5..afae79e3e 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 0cc1af7ba..622cc0ce0 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 90528637f..b2bafebb1 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}'