From 498059ba288a89c3a50958cd948f484d568d01e6 Mon Sep 17 00:00:00 2001
From: "James E. Blair" <jeblair@redhat.com>
Date: Tue, 20 Dec 2016 13:50:13 -0800
Subject: [PATCH] Add Zookeeper to tests

Add a requirement on kazoo and add a Zookeeper chroot to the test
infrastructure.

This is based on similar code in Nodepool.

Change-Id: Ic05386aac284c5542721fa3dcb1cd1c8e52d4a1f
---
 bindep.txt       |  1 +
 requirements.txt |  1 +
 tests/base.py    | 62 ++++++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 64 insertions(+)

diff --git a/bindep.txt b/bindep.txt
index a2cc02e43f..b0c4c3bf24 100644
--- a/bindep.txt
+++ b/bindep.txt
@@ -1 +1,2 @@
 libjpeg-dev [test]
+zookeeperd [platform:dpkg]
diff --git a/requirements.txt b/requirements.txt
index ec0b76ac48..4c5adc763b 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -16,3 +16,4 @@ PrettyTable>=0.6,<0.8
 babel>=1.0
 six>=1.6.0
 ansible>=2.0.0.1
+kazoo
diff --git a/tests/base.py b/tests/base.py
index 603fba3e8e..9845484ca0 100755
--- a/tests/base.py
+++ b/tests/base.py
@@ -1,6 +1,7 @@
 #!/usr/bin/env python
 
 # Copyright 2012 Hewlett-Packard Development Company, L.P.
+# Copyright 2016 Red Hat, Inc.
 #
 # Licensed under the Apache License, Version 2.0 (the "License"); you may
 # not use this file except in compliance with the License. You may obtain
@@ -39,6 +40,7 @@ import time
 import git
 import gear
 import fixtures
+import kazoo.client
 import statsd
 import testtools
 from git.exc import NoSuchPathError
@@ -869,6 +871,57 @@ class FakeSwiftClientConnection(swiftclient.client.Connection):
         return endpoint, ''
 
 
+class ChrootedKazooFixture(fixtures.Fixture):
+    def __init__(self):
+        super(ChrootedKazooFixture, self).__init__()
+
+        zk_host = os.environ.get('NODEPOOL_ZK_HOST', 'localhost')
+        if ':' in zk_host:
+            host, port = zk_host.split(':')
+        else:
+            host = zk_host
+            port = None
+
+        self.zookeeper_host = host
+
+        if not port:
+            self.zookeeper_port = 2181
+        else:
+            self.zookeeper_port = int(port)
+
+    def _setUp(self):
+        # Make sure the test chroot paths do not conflict
+        random_bits = ''.join(random.choice(string.ascii_lowercase +
+                                            string.ascii_uppercase)
+                              for x in range(8))
+
+        rand_test_path = '%s_%s' % (random_bits, os.getpid())
+        self.zookeeper_chroot = "/nodepool_test/%s" % rand_test_path
+
+        # Ensure the chroot path exists and clean up any 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.zookeeper_chroot):
+            _tmp_client.delete(self.zookeeper_chroot, recursive=True)
+
+        _tmp_client.ensure_path(self.zookeeper_chroot)
+        _tmp_client.stop()
+        _tmp_client.close()
+
+        self.addCleanup(self._cleanup)
+
+    def _cleanup(self):
+        '''Remove the chroot path.'''
+        # 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.zookeeper_chroot, recursive=True)
+        _tmp_client.stop()
+
+
 class BaseTestCase(testtools.TestCase):
     log = logging.getLogger("zuul.test")
 
@@ -989,6 +1042,9 @@ class ZuulTestCase(BaseTestCase):
 
     def setUp(self):
         super(ZuulTestCase, self).setUp()
+
+        self.setupZK()
+
         if USE_TEMPDIR:
             tmp_root = self.useFixture(fixtures.TempDir(
                 rootdir=os.environ.get("ZUUL_TEST_ROOT"))
@@ -1186,6 +1242,12 @@ class ZuulTestCase(BaseTestCase):
                     self.copyDirToRepo(project,
                                        os.path.join(git_path, reponame))
 
+    def setupZK(self):
+        self.zk_chroot_fixture = self.useFixture(ChrootedKazooFixture())
+        self.zookeeper_host = self.zk_chroot_fixture.zookeeper_host
+        self.zookeeper_port = self.zk_chroot_fixture.zookeeper_port
+        self.zookeeper_chroot = self.zk_chroot_fixture.zookeeper_chroot
+
     def copyDirToRepo(self, project, source_path):
         self.init_repo(project)