Clear out EPU-layer data in ZK setup step

And add tests

epu/test/test_zkcli.py

@@ -0,0 +1,83 @@
+import os
+import unittest
+import tempfile
+
+import yaml
+import mock
+from kazoo.exceptions import NoNodeException, NodeExistsException
+
+from epu import zkcli
+
+_ZK_CONFIG = {"server": {"zookeeper": {"hosts": "zk1,zk2", "path": "/base"}}}
+
+
+def _write_temp_yaml(o):
+    fd, path = tempfile.mkstemp()
+    f = os.fdopen(fd, 'w')
+    try:
+        yaml.dump(o, f)
+    except:
+        os.remove(path)
+        raise
+    finally:
+        f.close()
+    return path
+
+
+@mock.patch("epu.zkcli.KazooClient")
+class ZKCliTests(unittest.TestCase):
+
+    @classmethod
+    def setUpClass(cls):
+        cls.zk_config = _write_temp_yaml(_ZK_CONFIG)
+
+    @classmethod
+    def tearDownClass(cls):
+        if hasattr(cls, 'zk_config') and cls.zk_config:
+            os.remove(cls.zk_config)
+
+    def test_setup(self, kazoo_cls):
+        kazoo_cls.return_value = kazoo = mock.Mock()
+
+        zkcli.main(['--config', self.zk_config, "setup"])
+        kazoo.ensure_path.assert_called_once_with("/base")
+        kazoo.delete.assert_called_once_with("/base/provisioner/disabled")
+        kazoo.create.assert_called_once_with("/base/pd/system_boot", "",
+            makepath=True)
+
+        kazoo.reset_mock()
+        kazoo.delete.side_effect = NoNodeException()
+        zkcli.main(['--config', self.zk_config, "setup"])
+        kazoo.delete.assert_called_once_with("/base/provisioner/disabled")
+
+        kazoo.reset_mock()
+        kazoo.create.side_effect = NodeExistsException()
+        zkcli.main(['--config', self.zk_config, "setup"])
+        kazoo.create.assert_called_once_with("/base/pd/system_boot", "",
+            makepath=True)
+
+    def test_setup_clean(self, kazoo_cls):
+        kazoo_cls.return_value = kazoo = mock.Mock()
+
+        zkcli.main(['--config', self.zk_config, "setup", "--clean"])
+        kazoo.ensure_path.assert_called_once_with("/base")
+        kazoo.delete.assert_called_once_with("/base", recursive=True)
+        kazoo.create.assert_called_once_with("/base/pd/system_boot", "",
+            makepath=True)
+
+    def test_setup_clean_epu(self, kazoo_cls):
+        kazoo_cls.return_value = kazoo = mock.Mock()
+
+        zkcli.main(['--config', self.zk_config, "setup", "--clean-epu"])
+        kazoo.ensure_path.assert_called_once_with("/base")
+        self.assertEqual(kazoo.delete.call_count, 2)
+        kazoo.delete.assert_any_call("/base/provisioner", recursive=True)
+        kazoo.delete.assert_any_call("/base/epum", recursive=True)
+        kazoo.create.assert_called_once_with("/base/pd/system_boot", "",
+            makepath=True)
+
+    def test_destroy(self, kazoo_cls):
+        kazoo_cls.return_value = kazoo = mock.Mock()
+
+        zkcli.main(['--config', self.zk_config, "destroy"])
+        kazoo.delete.assert_called_once_with("/base", recursive=True)

epu/zkcli.py

@@ -6,10 +6,11 @@
 import os
 
 from kazoo.client import KazooClient
-from kazoo.exceptions import NodeExistsException
+from kazoo.exceptions import NodeExistsException, NoNodeException
 
 from epu import zkutil
 from epu.processdispatcher.store import ProcessDispatcherZooKeeperStore
+from epu.provisioner.store import ProvisionerZooKeeperStore
 
 
 class CliError(Exception):
@@ -60,6 +61,26 @@ def cmd_setup(zk_config, args):
         kazoo.create(system_boot_path, "", makepath=True)
     except NodeExistsException:
         pass
+
+    if not args.clean:
+        if args.clean_epu:
+            # clean out the EPU-layer storage as it is not needed between restarts
+
+            kazoo.delete("%s/provisioner" % path, recursive=True)
+            kazoo.delete("%s/epum" % path, recursive=True)
+
+        else:
+            # at the least, clear out the provisioner DISABLED flag if set.
+            # this allows new VMs to be launched.
+
+            # WARNING: provisioner path is hardcoded here!
+            prov_disabled_path = "%s/provisioner/%s" % (path, ProvisionerZooKeeperStore.DISABLED_PATH)
+            prov_disabled_path = os.path.normpath(prov_disabled_path)
+            try:
+                kazoo.delete(prov_disabled_path)
+            except NoNodeException:
+                pass
+
     kazoo.stop()
 
 
@@ -72,7 +93,7 @@ def cmd_destroy(zk_config, args):
     kazoo.stop()
 
 
-def main():
+def main(args=None):
     parser = argparse.ArgumentParser(description='EPU ZooKeeper management utility')
     parser.add_argument('--config', '-c', metavar="config.yml",
         type=argparse.FileType('r'), required=True)
@@ -83,14 +104,16 @@ def main():
     setup_parser.add_argument("path", help="ZooKeeper base path", nargs='?')
     setup_parser.add_argument("--clean", help="Remove existing contents of path",
         action="store_true")
+    setup_parser.add_argument("--clean-epu", help="Remove EPU-layer data",
+        action="store_true")
     setup_parser.set_defaults(func=cmd_setup)
 
     destroy_parser = subparsers.add_parser("destroy",
         description=cmd_destroy.__doc__, help="delete path")
     destroy_parser.add_argument("path", help="ZooKeeper base path", nargs='?')
     destroy_parser.set_defaults(func=cmd_destroy)
 
-    args = parser.parse_args()
+    args = parser.parse_args(args=args)
     config = yaml.load(args.config)
 
     if not zkutil.is_zookeeper_enabled(config):