ubuntu-advantage: Add new config module to support ubuntu-advantage-tools

ubuntu-advantage-tools is a package for enabling and disabling extended
support services such as Extended Security Maintenance (ESM), Canonical
Livepatch and FIPS certified PPAs. Simplify Ubuntu Advantage setup on
machines by allowing users to provide a list of ubuntu-advantage commands
in cloud-config.

2 files changed, 270 insertions, 50 deletions

diff --git a/cloudinit/config/tests/test_snap.py b/cloudinit/config/tests/test_snap.py
index cb1205e9..988e7f7c 100644
--- a/cloudinit/config/tests/test_snap.py
+++ b/cloudinit/config/tests/test_snap.py
@@ -4,8 +4,8 @@ import re
 from six import StringIO
 
 from cloudinit.config.cc_snap import (
-    ASSERTIONS_FILE, add_assertions, handle, prepend_snap_commands,
-    maybe_install_squashfuse, run_commands, schema)
+    ASSERTIONS_FILE, add_assertions, handle, maybe_install_squashfuse,
+    run_commands, schema)
 from cloudinit.config.schema import validate_cloudconfig_schema
 from cloudinit import util
 from cloudinit.tests.helpers import CiTestCase, mock, wrap_and_call
@@ -158,54 +158,6 @@ class TestAddAssertions(CiTestCase):
             util.load_file(compare_file), util.load_file(assert_file))
 
 
-class TestPrepentSnapCommands(CiTestCase):
-
-    with_logs = True
-
-    def test_prepend_snap_commands_errors_on_neither_string_nor_list(self):
-        """Raise an error for each command which is not a string or list."""
-        orig_commands = ['ls', 1, {'not': 'gonna work'}, ['snap', 'list']]
-        with self.assertRaises(TypeError) as context_manager:
-            prepend_snap_commands(orig_commands)
-        self.assertEqual(
-            "Invalid snap config. These commands are not a string or list:\n"
-            "1\n{'not': 'gonna work'}",
-            str(context_manager.exception))
-
-    def test_prepend_snap_commands_warns_on_non_snap_string_commands(self):
-        """Warn on each non-snap for commands of type string."""
-        orig_commands = ['ls', 'snap list', 'touch /blah', 'snap install x']
-        fixed_commands = prepend_snap_commands(orig_commands)
-        self.assertEqual(
-            'WARNING: Non-snap commands in snap config:\n'
-            'ls\ntouch /blah\n',
-            self.logs.getvalue())
-        self.assertEqual(orig_commands, fixed_commands)
-
-    def test_prepend_snap_commands_prepends_on_non_snap_list_commands(self):
-        """Prepend 'snap' for each non-snap command of type list."""
-        orig_commands = [['ls'], ['snap', 'list'], ['snapa', '/blah'],
-                         ['snap', 'install', 'x']]
-        expected = [['snap', 'ls'], ['snap', 'list'],
-                    ['snap', 'snapa', '/blah'],
-                    ['snap', 'install', 'x']]
-        fixed_commands = prepend_snap_commands(orig_commands)
-        self.assertEqual('', self.logs.getvalue())
-        self.assertEqual(expected, fixed_commands)
-
-    def test_prepend_snap_commands_removes_first_item_when_none(self):
-        """Remove the first element of a non-snap command when it is None."""
-        orig_commands = [[None, 'ls'], ['snap', 'list'],
-                         [None, 'touch', '/blah'],
-                         ['snap', 'install', 'x']]
-        expected = [['ls'], ['snap', 'list'],
-                    ['touch', '/blah'],
-                    ['snap', 'install', 'x']]
-        fixed_commands = prepend_snap_commands(orig_commands)
-        self.assertEqual('', self.logs.getvalue())
-        self.assertEqual(expected, fixed_commands)
-
-
 class TestRunCommands(CiTestCase):
 
     with_logs = True
diff --git a/cloudinit/config/tests/test_ubuntu_advantage.py b/cloudinit/config/tests/test_ubuntu_advantage.py
new file mode 100644
index 00000000..0eeadd43
--- /dev/null
+++ b/cloudinit/config/tests/test_ubuntu_advantage.py
@@ -0,0 +1,268 @@
+# This file is part of cloud-init. See LICENSE file for license information.
+
+import re
+from six import StringIO
+
+from cloudinit.config.cc_ubuntu_advantage import (
+    handle, maybe_install_ua_tools, run_commands, schema)
+from cloudinit.config.schema import validate_cloudconfig_schema
+from cloudinit import util
+from cloudinit.tests.helpers import CiTestCase, mock
+
+
+# Module path used in mocks
+MPATH = 'cloudinit.config.cc_ubuntu_advantage'
+
+
+class FakeCloud(object):
+    def __init__(self, distro):
+        self.distro = distro
+
+
+class TestRunCommands(CiTestCase):
+
+    with_logs = True
+
+    def setUp(self):
+        super(TestRunCommands, self).setUp()
+        self.tmp = self.tmp_dir()
+
+    @mock.patch('%s.util.subp' % MPATH)
+    def test_run_commands_on_empty_list(self, m_subp):
+        """When provided with an empty list, run_commands does nothing."""
+        run_commands([])
+        self.assertEqual('', self.logs.getvalue())
+        m_subp.assert_not_called()
+
+    def test_run_commands_on_non_list_or_dict(self):
+        """When provided an invalid type, run_commands raises an error."""
+        with self.assertRaises(TypeError) as context_manager:
+            run_commands(commands="I'm Not Valid")
+        self.assertEqual(
+            "commands parameter was not a list or dict: I'm Not Valid",
+            str(context_manager.exception))
+
+    def test_run_command_logs_commands_and_exit_codes_to_stderr(self):
+        """All exit codes are logged to stderr."""
+        outfile = self.tmp_path('output.log', dir=self.tmp)
+
+        cmd1 = 'echo "HI" >> %s' % outfile
+        cmd2 = 'bogus command'
+        cmd3 = 'echo "MOM" >> %s' % outfile
+        commands = [cmd1, cmd2, cmd3]
+
+        mock_path = '%s.sys.stderr' % MPATH
+        with mock.patch(mock_path, new_callable=StringIO) as m_stderr:
+            with self.assertRaises(RuntimeError) as context_manager:
+                run_commands(commands=commands)
+
+        self.assertIsNotNone(
+            re.search(r'bogus: (command )?not found',
+                      str(context_manager.exception)),
+            msg='Expected bogus command not found')
+        expected_stderr_log = '\n'.join([
+            'Begin run command: {cmd}'.format(cmd=cmd1),
+            'End run command: exit(0)',
+            'Begin run command: {cmd}'.format(cmd=cmd2),
+            'ERROR: End run command: exit(127)',
+            'Begin run command: {cmd}'.format(cmd=cmd3),
+            'End run command: exit(0)\n'])
+        self.assertEqual(expected_stderr_log, m_stderr.getvalue())
+
+    def test_run_command_as_lists(self):
+        """When commands are specified as a list, run them in order."""
+        outfile = self.tmp_path('output.log', dir=self.tmp)
+
+        cmd1 = 'echo "HI" >> %s' % outfile
+        cmd2 = 'echo "MOM" >> %s' % outfile
+        commands = [cmd1, cmd2]
+        with mock.patch('%s.sys.stderr' % MPATH, new_callable=StringIO):
+            run_commands(commands=commands)
+
+        self.assertIn(
+            'DEBUG: Running user-provided ubuntu-advantage commands',
+            self.logs.getvalue())
+        self.assertEqual('HI\nMOM\n', util.load_file(outfile))
+        self.assertIn(
+            'WARNING: Non-ubuntu-advantage commands in ubuntu-advantage'
+            ' config:',
+            self.logs.getvalue())
+
+    def test_run_command_dict_sorted_as_command_script(self):
+        """When commands are a dict, sort them and run."""
+        outfile = self.tmp_path('output.log', dir=self.tmp)
+        cmd1 = 'echo "HI" >> %s' % outfile
+        cmd2 = 'echo "MOM" >> %s' % outfile
+        commands = {'02': cmd1, '01': cmd2}
+        with mock.patch('%s.sys.stderr' % MPATH, new_callable=StringIO):
+            run_commands(commands=commands)
+
+        expected_messages = [
+            'DEBUG: Running user-provided ubuntu-advantage commands']
+        for message in expected_messages:
+            self.assertIn(message, self.logs.getvalue())
+        self.assertEqual('MOM\nHI\n', util.load_file(outfile))
+
+
+class TestSchema(CiTestCase):
+
+    with_logs = True
+
+    def test_schema_warns_on_ubuntu_advantage_not_as_dict(self):
+        """If ubuntu-advantage configuration is not a dict, emit a warning."""
+        validate_cloudconfig_schema({'ubuntu-advantage': 'wrong type'}, schema)
+        self.assertEqual(
+            "WARNING: Invalid config:\nubuntu-advantage: 'wrong type' is not"
+            " of type 'object'\n",
+            self.logs.getvalue())
+
+    @mock.patch('%s.run_commands' % MPATH)
+    def test_schema_disallows_unknown_keys(self, _):
+        """Unknown keys in ubuntu-advantage configuration emit warnings."""
+        validate_cloudconfig_schema(
+            {'ubuntu-advantage': {'commands': ['ls'], 'invalid-key': ''}},
+            schema)
+        self.assertIn(
+            'WARNING: Invalid config:\nubuntu-advantage: Additional properties'
+            " are not allowed ('invalid-key' was unexpected)",
+            self.logs.getvalue())
+
+    def test_warn_schema_requires_commands(self):
+        """Warn when ubuntu-advantage configuration lacks commands."""
+        validate_cloudconfig_schema(
+            {'ubuntu-advantage': {}}, schema)
+        self.assertEqual(
+            "WARNING: Invalid config:\nubuntu-advantage: 'commands' is a"
+            " required property\n",
+            self.logs.getvalue())
+
+    @mock.patch('%s.run_commands' % MPATH)
+    def test_warn_schema_commands_is_not_list_or_dict(self, _):
+        """Warn when ubuntu-advantage:commands config is not a list or dict."""
+        validate_cloudconfig_schema(
+            {'ubuntu-advantage': {'commands': 'broken'}}, schema)
+        self.assertEqual(
+            "WARNING: Invalid config:\nubuntu-advantage.commands: 'broken' is"
+            " not of type 'object', 'array'\n",
+            self.logs.getvalue())
+
+    @mock.patch('%s.run_commands' % MPATH)
+    def test_warn_schema_when_commands_is_empty(self, _):
+        """Emit warnings when ubuntu-advantage:commands is empty."""
+        validate_cloudconfig_schema(
+            {'ubuntu-advantage': {'commands': []}}, schema)
+        validate_cloudconfig_schema(
+            {'ubuntu-advantage': {'commands': {}}}, schema)
+        self.assertEqual(
+            "WARNING: Invalid config:\nubuntu-advantage.commands: [] is too"
+            " short\nWARNING: Invalid config:\nubuntu-advantage.commands: {}"
+            " does not have enough properties\n",
+            self.logs.getvalue())
+
+    @mock.patch('%s.run_commands' % MPATH)
+    def test_schema_when_commands_are_list_or_dict(self, _):
+        """No warnings when ubuntu-advantage:commands are a list or dict."""
+        validate_cloudconfig_schema(
+            {'ubuntu-advantage': {'commands': ['valid']}}, schema)
+        validate_cloudconfig_schema(
+            {'ubuntu-advantage': {'commands': {'01': 'also valid'}}}, schema)
+        self.assertEqual('', self.logs.getvalue())
+
+
+class TestHandle(CiTestCase):
+
+    with_logs = True
+
+    def setUp(self):
+        super(TestHandle, self).setUp()
+        self.tmp = self.tmp_dir()
+
+    @mock.patch('%s.run_commands' % MPATH)
+    @mock.patch('%s.validate_cloudconfig_schema' % MPATH)
+    def test_handle_no_config(self, m_schema, m_run):
+        """When no ua-related configuration is provided, nothing happens."""
+        cfg = {}
+        handle('ua-test', cfg=cfg, cloud=None, log=self.logger, args=None)
+        self.assertIn(
+            "DEBUG: Skipping module named ua-test, no 'ubuntu-advantage' key"
+            " in config",
+            self.logs.getvalue())
+        m_schema.assert_not_called()
+        m_run.assert_not_called()
+
+    @mock.patch('%s.maybe_install_ua_tools' % MPATH)
+    def test_handle_tries_to_install_ubuntu_advantage_tools(self, m_install):
+        """If ubuntu_advantage is provided, try installing ua-tools package."""
+        cfg = {'ubuntu-advantage': {}}
+        mycloud = FakeCloud(None)
+        handle('nomatter', cfg=cfg, cloud=mycloud, log=self.logger, args=None)
+        m_install.assert_called_once_with(mycloud)
+
+    @mock.patch('%s.maybe_install_ua_tools' % MPATH)
+    def test_handle_runs_commands_provided(self, m_install):
+        """When commands are specified as a list, run them."""
+        outfile = self.tmp_path('output.log', dir=self.tmp)
+
+        cfg = {
+            'ubuntu-advantage': {'commands': ['echo "HI" >> %s' % outfile,
+                                              'echo "MOM" >> %s' % outfile]}}
+        mock_path = '%s.sys.stderr' % MPATH
+        with mock.patch(mock_path, new_callable=StringIO):
+            handle('nomatter', cfg=cfg, cloud=None, log=self.logger, args=None)
+        self.assertEqual('HI\nMOM\n', util.load_file(outfile))
+
+
+class TestMaybeInstallUATools(CiTestCase):
+
+    with_logs = True
+
+    def setUp(self):
+        super(TestMaybeInstallUATools, self).setUp()
+        self.tmp = self.tmp_dir()
+
+    @mock.patch('%s.util.which' % MPATH)
+    def test_maybe_install_ua_tools_noop_when_ua_tools_present(self, m_which):
+        """Do nothing if ubuntu-advantage-tools already exists."""
+        m_which.return_value = '/usr/bin/ubuntu-advantage'  # already installed
+        distro = mock.MagicMock()
+        distro.update_package_sources.side_effect = RuntimeError(
+            'Some apt error')
+        maybe_install_ua_tools(cloud=FakeCloud(distro))  # No RuntimeError
+
+    @mock.patch('%s.util.which' % MPATH)
+    def test_maybe_install_ua_tools_raises_update_errors(self, m_which):
+        """maybe_install_ua_tools logs and raises apt update errors."""
+        m_which.return_value = None
+        distro = mock.MagicMock()
+        distro.update_package_sources.side_effect = RuntimeError(
+            'Some apt error')
+        with self.assertRaises(RuntimeError) as context_manager:
+            maybe_install_ua_tools(cloud=FakeCloud(distro))
+        self.assertEqual('Some apt error', str(context_manager.exception))
+        self.assertIn('Package update failed\nTraceback', self.logs.getvalue())
+
+    @mock.patch('%s.util.which' % MPATH)
+    def test_maybe_install_ua_raises_install_errors(self, m_which):
+        """maybe_install_ua_tools logs and raises package install errors."""
+        m_which.return_value = None
+        distro = mock.MagicMock()
+        distro.update_package_sources.return_value = None
+        distro.install_packages.side_effect = RuntimeError(
+            'Some install error')
+        with self.assertRaises(RuntimeError) as context_manager:
+            maybe_install_ua_tools(cloud=FakeCloud(distro))
+        self.assertEqual('Some install error', str(context_manager.exception))
+        self.assertIn(
+            'Failed to install ubuntu-advantage-tools\n', self.logs.getvalue())
+
+    @mock.patch('%s.util.which' % MPATH)
+    def test_maybe_install_ua_tools_happy_path(self, m_which):
+        """maybe_install_ua_tools installs ubuntu-advantage-tools."""
+        m_which.return_value = None
+        distro = mock.MagicMock()  # No errors raised
+        maybe_install_ua_tools(cloud=FakeCloud(distro))
+        distro.update_package_sources.assert_called_once_with()
+        distro.install_packages.assert_called_once_with(
+            ['ubuntu-advantage-tools'])
+
+# vi: ts=4 expandtab