implement power_state with tests.

4 files changed, 175 insertions, 78 deletions

diff --git a/ChangeLog b/ChangeLog
index 57bfc1a4..23612228 100644
--- a/ChangeLog
+++ b/ChangeLog
@@ -43,8 +43,8 @@
     - Added dependency on distribute's python-pkg-resources
  - use a set of helper/parsing classes to perform system configuration
    for easier test.  (/etc/sysconfig, /etc/hostname, resolv.conf, /etc/hosts)
- - add 'finalcmd' config module to execute 'finalcmd' entries like
-   'runcmd' but detached from cloud-init (LP: #1064665)
+ - add power_state_change config module for shutting down stystem after
+   cloud-init finishes. (LP: #1064665)
 0.7.0:
  - add a 'exception_cb' argument to 'wait_for_url'.  If provided, this
    method will be called back with the exception received and the message.
diff --git a/cloudinit/config/cc_power_state_change.py b/cloudinit/config/cc_power_state_change.py
index 67e0316b..07de548c 100644
--- a/cloudinit/config/cc_power_state_change.py
+++ b/cloudinit/config/cc_power_state_change.py
@@ -21,6 +21,7 @@ from cloudinit import util
 
 import errno
 import os
+import re
 import subprocess
 import sys
 import time
@@ -28,83 +29,91 @@ import time
 frequency = PER_INSTANCE
 
 
-def handle(_name, cfg, cloud, log, _args):
+def handle(_name, cfg, _cloud, log, _args):
 
-    finalcmds = cfg.get("finalcmd")
-
-    if not finalcmds:
-        log.debug("No final commands")
+    try:
+        (args, timeout) = load_power_state(cfg)
+        if args is None:
+            log.debug("no power_state provided. doing nothing")
+            return
+    except Exception as e:
+        log.warn("%s Not performing power state change!" % str(e))
         return
 
     mypid = os.getpid()
     cmdline = util.load_file("/proc/%s/cmdline")
 
     if not cmdline:
-        log.warn("Failed to get cmdline of current process")
-        return
-
-    try:
-        timeout = float(cfg.get("finalcmd_timeout", 30.0))
-    except ValueError:
-        log.warn("failed to convert finalcmd_timeout '%s' to float" %
-                 cfg.get("finalcmd_timeout", 30.0))
+        log.warn("power_state: failed to get cmdline of current process")
         return
 
     devnull_fp = open("/dev/null", "w")
 
-    shellcode = util.shellify(finalcmds)
+    log.debug("After pid %s ends, will execute: %s" % (mypid, ' '.join(args)))
+
+    util.fork_cb(run_after_pid_gone, mypid, cmdline, timeout, log, execmd,
+                 [args, devnull_fp])
+
+
+def load_power_state(cfg):
+    # returns a tuple of shutdown_command, timeout
+    # shutdown_command is None if no config found
+    pstate = cfg.get('power_state')
+
+    if pstate is None:
+        return (None, None)
+
+    if not isinstance(pstate, dict):
+        raise TypeError("power_state is not a dict.")
 
-    # note, after the fork, we do not use any of cloud-init's functions
-    # that would attempt to log.  The primary reason for that is
-    # to allow the 'finalcmd' the ability to do just about anything
-    # and not depend on syslog services.
-    # Basically, it should "just work" to have finalcmd of:
-    #  - sleep 30
-    #  - /sbin/poweroff
-    finalcmd_d = os.path.join(cloud.get_ipath_cur(), "finalcmds")
+    opt_map = {'halt': '-H', 'poweroff': '-P', 'reboot': '-r'}
 
-    util.fork_cb(run_after_pid_gone, mypid, cmdline, timeout,
-                 runfinal, (shellcode, finalcmd_d, devnull_fp))
+    mode = pstate.get("mode")
+    if mode not in opt_map:
+        raise TypeError("power_state[mode] required, must be one of: %s." %
+                        ','.join(opt_map.keys()))
 
+    delay = pstate.get("delay", "now")
+    if delay != "now" and not re.match("\+[0-9]+", delay):
+        raise TypeError("power_state[delay] must be 'now' or '+m' (minutes).")
 
-def execmd(exe_args, data_in=None, output=None):
+    args = ["shutdown", opt_map[mode], delay]
+    if pstate.get("message"):
+        args.append(pstate.get("message"))
+
+    try:
+        timeout = float(pstate.get('timeout', 30.0))
+    except ValueError:
+        raise ValueError("failed to convert timeout '%s' to float." %
+                         pstate['timeout'])
+
+    return (args, timeout)
+
+
+def execmd(exe_args, output=None, data_in=None):
     try:
         proc = subprocess.Popen(exe_args, stdin=subprocess.PIPE,
                                 stdout=output, stderr=subprocess.STDOUT)
         proc.communicate(data_in)
     except Exception:
-        return 254
-    return proc.returncode()
+        sys.exit(254)
+    sys.exit(proc.returncode())
 
 
-def runfinal(shellcode, finalcmd_d, output=None):
-    ret = execmd(["/bin/sh"], data_in=shellcode, output=output)
-    if not (finalcmd_d and os.path.isdir(finalcmd_d)):
-        sys.exit(ret)
-
-    fails = 0
-    if ret != 0:
-        fails = 1
-
-    # now runparts the final command dir
-    for exe_name in sorted(os.listdir(finalcmd_d)):
-        exe_path = os.path.join(finalcmd_d, exe_name)
-        if os.path.isfile(exe_path) and os.access(exe_path, os.X_OK):
-            ret = execmd([exe_path], data_in=None, output=output)
-            if ret != 0:
-                fails += 1
-    sys.exit(fails)
-
-
-def run_after_pid_gone(pid, pidcmdline, timeout, func, args):
+def run_after_pid_gone(pid, pidcmdline, timeout, log, func, args):
     # wait until pid, with /proc/pid/cmdline contents of pidcmdline
     # is no longer alive.  After it is gone, or timeout has passed
     # execute func(args)
-    msg = "ERROR: Uncaught error"
+    msg = None
     end_time = time.time() + timeout
 
     cmdline_f = "/proc/%s/cmdline" % pid
 
+    def fatal(msg):
+        if log:
+            log.warn(msg)
+        sys.exit(254)
+
     while True:
         if time.time() > end_time:
             msg = "timeout reached before %s ended" % pid
@@ -122,18 +131,15 @@ def run_after_pid_gone(pid, pidcmdline, timeout, func, args):
             if ioerr.errno == errno.ENOENT:
                 msg = "pidfile '%s' gone" % cmdline_f
             else:
-                msg = "ERROR: IOError: %s" % ioerr
-                raise
+                fatal("IOError during wait: %s" % ioerr)
             break
 
         except Exception as e:
-            msg = "ERROR: Exception: %s" % e
-            raise
+            fatal("Unexpected Exception: %s" % e)
 
-    if msg.startswith("ERROR:"):
-        sys.stderr.write(msg)
-        sys.stderr.write("Not executing finalcmd")
-        sys.exit(1)
+    if not msg:
+        fatal("Unexpected error in run_after_pid_gone")
 
-    sys.stderr.write("calling %s with %s\n" % (func, args))
-    sys.exit(func(*args))
+    if log:
+        log.debug(msg)
+    func(*args)
diff --git a/doc/examples/cloud-config.txt b/doc/examples/cloud-config.txt
index 4fc5f351..09298655 100644
--- a/doc/examples/cloud-config.txt
+++ b/doc/examples/cloud-config.txt
@@ -256,24 +256,6 @@ bootcmd:
  - echo 192.168.1.130 us.archive.ubuntu.com > /etc/hosts
  - [ cloud-init-per, once, mymkfs, mkfs, /dev/vdb ]
 
-# final commands
-# default: none
-# This can be used to execute commands after and fully detached from
-# a cloud-init stage.  The initial purpose of it was to allow 'poweroff'
-# detached from cloud-init.  If poweroff was run from 'runcmd' or userdata
-# then messages may be spewed from cloud-init about logging failing or other
-# issues as a result of the system being turned off.
-#
-# You probably are better off using 'runcmd' for this.
-#
-# The output of finalcmd will redirected redirected to /dev/null
-# If you want output to be seen, take care to do so in your commands
-# themselves.  See example.
-finalcmd:
- - sleep 30
- - "echo $(date -R): powering off > /dev/console"
- - /sbin/poweroff
-
 # cloud_config_modules:
 # default:
 # cloud_config_modules:
@@ -596,3 +578,24 @@ manual_cache_clean: False
 #   A list of key types (first token of a /etc/ssh/ssh_key_*.pub file)
 #   that should be skipped when outputting key fingerprints and keys
 #   to the console respectively.
+
+## poweroff or reboot system after finished
+# default: none
+#
+# power_state can be used to make the system shutdown, reboot or
+# halt after boot is finished.  This same thing can be acheived by
+# user-data scripts or by runcmd by simply invoking 'shutdown'.
+# 
+# Doing it this way ensures that cloud-init is entirely finished with
+# modules that would be executed, and avoids any error/log messages
+# that may go to the console as a result of system services like
+# syslog being taken down while cloud-init is running.
+#
+# delay: form accepted by shutdown.  default is 'now'. other format
+#        accepted is +m (m in minutes)
+# mode: required. must be one of 'poweroff', 'halt', 'reboot'
+# message: provided as the message argument to 'shutdown'. default is none.
+power_state:
+ delay: 30
+ mode: poweroff
+ message: Bye Bye
diff --git a/tests/unittests/test_handler/test_handler_power_state.py b/tests/unittests/test_handler/test_handler_power_state.py
new file mode 100644
index 00000000..1149fedc
--- /dev/null
+++ b/tests/unittests/test_handler/test_handler_power_state.py
@@ -0,0 +1,88 @@
+from unittest import TestCase
+
+from cloudinit.config import cc_power_state_change as psc
+
+
+class TestLoadPowerState(TestCase):
+    def setUp(self):
+        super(self.__class__, self).setUp()
+
+    def test_no_config(self):
+        # completely empty config should mean do nothing
+        (cmd, _timeout) = psc.load_power_state({})
+        self.assertEqual(cmd, None)
+
+    def test_irrelevant_config(self):
+        # no power_state field in config should return None for cmd
+        (cmd, _timeout) = psc.load_power_state({'foo': 'bar'})
+        self.assertEqual(cmd, None)
+
+    def test_invalid_mode(self):
+        cfg = {'power_state': {'mode': 'gibberish'}}
+        self.assertRaises(TypeError, psc.load_power_state, cfg)
+
+        cfg = {'power_state': {'mode': ''}}
+        self.assertRaises(TypeError, psc.load_power_state, cfg)
+
+    def test_empty_mode(self):
+        cfg = {'power_state': {'message': 'goodbye'}}
+        self.assertRaises(TypeError, psc.load_power_state, cfg)
+
+    def test_valid_modes(self):
+        cfg = {'power_state': {}}
+        for mode in ('halt', 'poweroff', 'reboot'):
+            cfg['power_state']['mode'] = mode
+            check_lps_ret(psc.load_power_state(cfg), mode=mode)
+
+    def test_invalid_delay(self):
+        cfg = {'power_state': {'mode': 'poweroff', 'delay': 'goodbye'}}
+        self.assertRaises(TypeError, psc.load_power_state, cfg)
+
+    def test_valid_delay(self):
+        cfg = {'power_state': {'mode': 'poweroff', 'delay': ''}}
+        for delay in ("now", "+1", "+30"):
+            cfg['power_state']['delay'] = delay
+            check_lps_ret(psc.load_power_state(cfg))
+
+    def test_message_present(self):
+        cfg = {'power_state': {'mode': 'poweroff', 'message': 'GOODBYE'}}
+        ret = psc.load_power_state(cfg)
+        check_lps_ret(psc.load_power_state(cfg))
+        self.assertIn(cfg['power_state']['message'], ret[0])
+
+    def test_no_message(self):
+        # if message is not present, then no argument should be passed for it
+        cfg = {'power_state': {'mode': 'poweroff'}}
+        (cmd, _timeout) = psc.load_power_state(cfg)
+        self.assertNotIn("", cmd)
+        check_lps_ret(psc.load_power_state(cfg))
+        self.assertTrue(len(cmd) == 3)
+
+
+def check_lps_ret(psc_return, mode=None):
+    if len(psc_return) != 2:
+        raise TypeError("length returned = %d" % len(psc_return))
+
+    errs = []
+    cmd = psc_return[0]
+    timeout = psc_return[1]
+
+    if not 'shutdown' in psc_return[0][0]:
+        errs.append("string 'shutdown' not in cmd")
+
+    if mode is not None:
+        opt = {'halt': '-H', 'poweroff': '-P', 'reboot': '-r'}[mode]
+        if opt not in psc_return[0]:
+            errs.append("opt '%s' not in cmd: %s" % (opt, cmd))
+
+    if len(cmd) != 3 and len(cmd) != 4:
+        errs.append("Invalid command length: %s" % len(cmd))
+
+    try:
+        float(timeout)
+    except:
+        errs.append("timeout failed convert to float")
+
+    if len(errs):
+        lines = ["Errors in result: %s" % str(psc_return)] + errs
+        raise Exception('\n'.join(lines))