Merge lp:~milner/cloud-init/protect-from-bad-part-handlers into lp:~cloud-init-dev/cloud-init/trunk

[Mike Milner]

Catch exceptions from part-handlers and log the error before continuing.
This branch also adds tests for part-handler registration and part-handler handling.

=== modified file 'cloudinit/__init__.py'
--- cloudinit/__init__.py	2012-02-16 18:31:19 +0000
+++ cloudinit/__init__.py	2012-02-21 19:55:32 +0000
@@ -60,7 +60,6 @@
 import sys
 import os.path
 import errno
-import pwd
 import subprocess
 import yaml
 import logging
@@ -572,10 +571,14 @@
     if not (modfreq == per_always or
             (frequency == per_instance and modfreq == per_instance)):
         return
-    if mod.handler_version == 1:
-        mod.handle_part(data, ctype, filename, payload)
-    else:
-        mod.handle_part(data, ctype, filename, payload, frequency)
+    try:
+        if mod.handler_version == 1:
+            mod.handle_part(data, ctype, filename, payload)
+        else:
+            mod.handle_part(data, ctype, filename, payload, frequency)
+    except:
+        util.logexc(log)
+        traceback.print_exc(file=sys.stderr)
 
 
 def partwalker_handle_handler(pdata, _ctype, _filename, payload):
@@ -586,15 +589,13 @@
     modfname = modname + ".py"
     util.write_file("%s/%s" % (pdata['handlerdir'], modfname), payload, 0600)
 
-    pdata['handlercount'] = curcount + 1
-
     try:
         mod = __import__(modname)
         handler_register(mod, pdata['handlers'], pdata['data'], frequency)
+        pdata['handlercount'] = curcount + 1
     except:
         util.logexc(log)
         traceback.print_exc(file=sys.stderr)
-        return
 
 
 def partwalker_callback(pdata, ctype, filename, payload):
 
=== added file 'tests/unittests/test__init__.py'
--- tests/unittests/test__init__.py	1970-01-01 00:00:00 +0000
+++ tests/unittests/test__init__.py	2012-02-21 19:55:32 +0000
@@ -0,0 +1,195 @@
+from mocker import MockerTestCase, ANY, ARGS, KWARGS
+import os
+
+from cloudinit import (partwalker_handle_handler, handler_handle_part,
+                       handler_register)
+from cloudinit.util import write_file, logexc
+
+
+class TestPartwalkerHandleHandler(MockerTestCase):
+    def setUp(self):
+        self.data = {
+            "handlercount": 0,
+            "frequency": "?",
+            "handlerdir": "?",
+            "handlers": [],
+            "data": None}
+
+        self.expected_module_name = "part-handler-%03d" % (
+            self.data["handlercount"],)
+        expected_file_name = "%s.py" % self.expected_module_name
+        expected_file_fullname = os.path.join(self.data["handlerdir"],
+                                              expected_file_name)
+        self.module_fake = "fake module handle"
+        self.ctype = None
+        self.filename = None
+        self.payload = "dummy payload"
+
+        # Mock the write_file function
+        write_file_mock = self.mocker.replace(write_file, passthrough=False)
+        write_file_mock(expected_file_fullname, self.payload, 0600)
+
+    def test_no_errors(self):
+        """Payload gets written to file and added to C{pdata}."""
+        # Mock the __import__ builtin
+        import_mock = self.mocker.replace("__builtin__.__import__")
+        import_mock(self.expected_module_name)
+        self.mocker.result(self.module_fake)
+        # Mock the handle_register function
+        handle_reg_mock = self.mocker.replace(handler_register,
+                                              passthrough=False)
+        handle_reg_mock(self.module_fake, self.data["handlers"],
+                        self.data["data"], self.data["frequency"])
+        # Activate mocks
+        self.mocker.replay()
+
+        partwalker_handle_handler(self.data, self.ctype, self.filename,
+                                  self.payload)
+
+        self.assertEqual(1, self.data["handlercount"])
+
+    def test_import_error(self):
+        """Module import errors are logged. No handler added to C{pdata}"""
+        # Mock the __import__ builtin
+        import_mock = self.mocker.replace("__builtin__.__import__")
+        import_mock(self.expected_module_name)
+        self.mocker.throw(ImportError())
+        # Mock log function
+        logexc_mock = self.mocker.replace(logexc, passthrough=False)
+        logexc_mock(ANY)
+        # Mock the print_exc function
+        print_exc_mock = self.mocker.replace("traceback.print_exc",
+                                             passthrough=False)
+        print_exc_mock(ARGS, KWARGS)
+        # Activate mocks
+        self.mocker.replay()
+
+        partwalker_handle_handler(self.data, self.ctype, self.filename,
+                                  self.payload)
+
+        self.assertEqual(0, self.data["handlercount"])
+
+    def test_attribute_error(self):
+        """Attribute errors are logged. No handler added to C{pdata}"""
+        # Mock the __import__ builtin
+        import_mock = self.mocker.replace("__builtin__.__import__")
+        import_mock(self.expected_module_name)
+        self.mocker.result(self.module_fake)
+        # Mock the handle_register function
+        handle_reg_mock = self.mocker.replace(handler_register,
+                                              passthrough=False)
+        handle_reg_mock(self.module_fake, self.data["handlers"],
+                        self.data["data"], self.data["frequency"])
+        self.mocker.throw(AttributeError())
+        # Mock log function
+        logexc_mock = self.mocker.replace(logexc, passthrough=False)
+        logexc_mock(ANY)
+        # Mock the print_exc function
+        print_exc_mock = self.mocker.replace("traceback.print_exc",
+                                             passthrough=False)
+        print_exc_mock(ARGS, KWARGS)
+        # Activate mocks
+        self.mocker.replay()
+
+        partwalker_handle_handler(self.data, self.ctype, self.filename,
+                                  self.payload)
+
+        self.assertEqual(0, self.data["handlercount"])
+
+
+class TestHandlerHandlePart(MockerTestCase):
+    def setUp(self):
+        self.data = "fake data"
+        self.ctype = "fake ctype"
+        self.filename = "fake filename"
+        self.payload = "fake payload"
+        self.frequency = "once-per-instance"
+
+    def test_normal_version_1(self):
+        """
+        C{handle_part} is called without C{frequency} for
+        C{handler_version} == 1.
+        """
+        # Build a mock part-handler module
+        mod_mock = self.mocker.mock()
+        getattr(mod_mock, "frequency")
+        self.mocker.result("once-per-instance")
+        getattr(mod_mock, "handler_version")
+        self.mocker.result(1)
+        mod_mock.handle_part(self.data, self.ctype, self.filename,
+                             self.payload)
+        self.mocker.replay()
+
+        handler_handle_part(mod_mock, self.data, self.ctype, self.filename,
+                            self.payload, self.frequency)
+
+    def test_normal_version_2(self):
+        """
+        C{handle_part} is called with C{frequency} for
+        C{handler_version} == 2.
+        """
+        # Build a mock part-handler module
+        mod_mock = self.mocker.mock()
+        getattr(mod_mock, "frequency")
+        self.mocker.result("once-per-instance")
+        getattr(mod_mock, "handler_version")
+        self.mocker.result(2)
+        mod_mock.handle_part(self.data, self.ctype, self.filename,
+                             self.payload, self.frequency)
+        self.mocker.replay()
+
+        handler_handle_part(mod_mock, self.data, self.ctype, self.filename,
+                            self.payload, self.frequency)
+
+    def test_modfreq_per_always(self):
+        """
+        C{handle_part} is called regardless of frequency if nofreq is always.
+        """
+        self.frequency = "once"
+        # Build a mock part-handler module
+        mod_mock = self.mocker.mock()
+        getattr(mod_mock, "frequency")
+        self.mocker.result("always")
+        getattr(mod_mock, "handler_version")
+        self.mocker.result(1)
+        mod_mock.handle_part(self.data, self.ctype, self.filename,
+                             self.payload)
+        self.mocker.replay()
+
+        handler_handle_part(mod_mock, self.data, self.ctype, self.filename,
+                            self.payload, self.frequency)
+
+    def test_no_handle_when_modfreq_once(self):
+        """C{handle_part} is not called if frequency is once"""
+        self.frequency = "once"
+        # Build a mock part-handler module
+        mod_mock = self.mocker.mock()
+        getattr(mod_mock, "frequency")
+        self.mocker.result("once-per-instance")
+        self.mocker.replay()
+
+        handler_handle_part(mod_mock, self.data, self.ctype, self.filename,
+                            self.payload, self.frequency)
+
+    def test_exception_is_caught(self):
+        """Exceptions within C{handle_part} are caught and logged."""
+        # Build a mock part-handler module
+        mod_mock = self.mocker.mock()
+        getattr(mod_mock, "frequency")
+        self.mocker.result("once-per-instance")
+        getattr(mod_mock, "handler_version")
+        self.mocker.result(1)
+        mod_mock.handle_part(self.data, self.ctype, self.filename,
+                             self.payload)
+        self.mocker.throw(Exception())
+        # Mock log function
+        logexc_mock = self.mocker.replace(logexc, passthrough=False)
+        logexc_mock(ANY)
+        # Mock the print_exc function
+        print_exc_mock = self.mocker.replace("traceback.print_exc",
+                                             passthrough=False)
+        print_exc_mock(ARGS, KWARGS)
+        self.mocker.replay()
+
+        handler_handle_part(mod_mock, self.data, self.ctype, self.filename,
+                            self.payload, self.frequency)
 
=== added directory 'tests/unittests/test_handler'
=== renamed file 'tests/unittests/test_handler_ca_certs.py' => 'tests/unittests/test_handler/test_handler_ca_certs.py'