cc_puppet: support AIO installations and more (#960)

- update the puppet module to support AIO installations by setting
  `install_type` to `aio`
- make the install collection configurable through the `collection`
  parameter; by default the rolling `puppet` collection will be used,
  which installs the latest version)
- when `install_type` is `aio`, puppetlabs repos will be purged after
  installation; set `cleanup` to `False` to prevent this
- AIO installations are performed by downloading and executing a shell
  script; the URL for this script can be overridden using the
  `aio_install_url` parameter
- make it possible to run puppet agent after installation/configuration
  via the `exec` key
- by default, puppet agent will run with the `--test` argument; this can
  be overridden via the `exec_args` key

4 files changed, 426 insertions, 64 deletions

diff --git a/cloudinit/config/cc_puppet.py b/cloudinit/config/cc_puppet.py
index bc981cf4..a0779eb0 100644
--- a/cloudinit/config/cc_puppet.py
+++ b/cloudinit/config/cc_puppet.py
@@ -29,22 +29,41 @@ The keys are ``package_name``, ``conf_file``, ``ssl_dir`` and
 ones that work with puppet 3.x and with distributions that ship modified
 puppet 4.x that uses the old paths.
 
+Agent packages from the puppetlabs repositories can be installed by setting
+``install_type`` to ``aio``. Based on this setting, the default config/SSL/CSR
+paths will be adjusted accordingly. To maintain backwards compatibility this
+setting defaults to ``packages`` which will install puppet from the distro
+packages.
+
+If installing ``aio`` packages, ``collection`` can also be set to one of
+``puppet`` (rolling release), ``puppet6``, ``puppet7`` (or their nightly
+counterparts) in order to install specific release streams. By default, the
+puppetlabs repository will be purged after installation finishes; set
+``cleanup`` to ``false`` to prevent this. AIO packages are installed through a
+shell script which is downloaded on the machine and then executed; the path to
+this script can be overridden using the ``aio_install_url`` key.
+
 Puppet configuration can be specified under the ``conf`` key. The
 configuration is specified as a dictionary containing high-level ``<section>``
 keys and lists of ``<key>=<value>`` pairs within each section. Each section
 name and ``<key>=<value>`` pair is written directly to ``puppet.conf``. As
-such,  section names should be one of: ``main``, ``master``, ``agent`` or
+such,  section names should be one of: ``main``, ``server``, ``agent`` or
 ``user`` and keys should be valid puppet configuration options. The
 ``certname`` key supports string substitutions for ``%i`` and ``%f``,
 corresponding to the instance id and fqdn of the machine respectively.
 If ``ca_cert`` is present, it will not be written to ``puppet.conf``, but
-instead will be used as the puppermaster certificate. It should be specified
+instead will be used as the puppetserver certificate. It should be specified
 in pem format as a multi-line string (using the ``|`` yaml notation).
 
-Additionally it's possible to create a csr_attributes.yaml for
-CSR attributes and certificate extension requests.
+Additionally it's possible to create a ``csr_attributes.yaml`` file for CSR
+attributes and certificate extension requests.
 See https://puppet.com/docs/puppet/latest/config_file_csr_attributes.html
 
+The puppet service will be automatically enabled after installation. A manual
+run can also be triggered by setting ``exec`` to ``true``, and additional
+arguments can be passed to ``puppet agent`` via the ``exec_args`` key (by
+default the agent will execute with the ``--test`` flag).
+
 **Internal name:** ``cc_puppet``
 
 **Module frequency:** per instance
@@ -56,13 +75,19 @@ See https://puppet.com/docs/puppet/latest/config_file_csr_attributes.html
     puppet:
         install: <true/false>
         version: <version>
+        collection: <aio collection>
+        install_type: <packages/aio>
+        aio_install_url: 'https://git.io/JBhoQ'
+        cleanup: <true/false>
         conf_file: '/etc/puppet/puppet.conf'
         ssl_dir: '/var/lib/puppet/ssl'
         csr_attributes_path: '/etc/puppet/csr_attributes.yaml'
         package_name: 'puppet'
+        exec: <true/false>
+        exec_args: ['--test']
         conf:
             agent:
-                server: "puppetmaster.example.org"
+                server: "puppetserver.example.org"
                 certname: "%i.%f"
                 ca_cert: |
                     -------BEGIN CERTIFICATE-------
@@ -84,12 +109,12 @@ from io import StringIO
 
 from cloudinit import helpers
 from cloudinit import subp
+from cloudinit import temp_utils
 from cloudinit import util
+from cloudinit import url_helper
 
-PUPPET_CONF_PATH = '/etc/puppet/puppet.conf'
-PUPPET_SSL_DIR = '/var/lib/puppet/ssl'
-PUPPET_CSR_ATTRIBUTES_PATH = '/etc/puppet/csr_attributes.yaml'
-PUPPET_PACKAGE_NAME = 'puppet'
+AIO_INSTALL_URL = 'https://raw.githubusercontent.com/puppetlabs/install-puppet/main/install.sh'  # noqa: E501
+PUPPET_AGENT_DEFAULT_ARGS = ['--test']
 
 
 class PuppetConstants(object):
@@ -119,6 +144,43 @@ def _autostart_puppet(log):
                      " puppet services on this system"))
 
 
+def get_config_value(puppet_bin, setting):
+    """Get the config value for a given setting using `puppet config print`
+    :param puppet_bin: path to puppet binary
+    :param setting: setting to query
+    """
+    out, _ = subp.subp([puppet_bin, 'config', 'print', setting])
+    return out.rstrip()
+
+
+def install_puppet_aio(url=AIO_INSTALL_URL, version=None,
+                       collection=None, cleanup=True):
+    """Install puppet-agent from the puppetlabs repositories using the one-shot
+    shell script
+
+    :param url: URL from where to download the install script
+    :param version: version to install, blank defaults to latest
+    :param collection: collection to install, blank defaults to latest
+    :param cleanup: whether to purge the puppetlabs repo after installation
+    """
+    args = []
+    if version is not None:
+        args = ['-v', version]
+    if collection is not None:
+        args += ['-c', collection]
+
+    # Purge puppetlabs repos after installation
+    if cleanup:
+        args += ['--cleanup']
+    content = url_helper.readurl(url=url, retries=5).contents
+
+    # Use tmpdir over tmpfile to avoid 'text file busy' on execute
+    with temp_utils.tempdir(needs_exe=True) as tmpd:
+        tmpf = os.path.join(tmpd, 'puppet-install')
+        util.write_file(tmpf, content, mode=0o700)
+        return subp.subp([tmpf] + args, capture=False)
+
+
 def handle(name, cfg, cloud, log, _args):
     # If there isn't a puppet key in the configuration don't do anything
     if 'puppet' not in cfg:
@@ -130,23 +192,50 @@ def handle(name, cfg, cloud, log, _args):
     # Start by installing the puppet package if necessary...
     install = util.get_cfg_option_bool(puppet_cfg, 'install', True)
     version = util.get_cfg_option_str(puppet_cfg, 'version', None)
-    package_name = util.get_cfg_option_str(
-        puppet_cfg, 'package_name', PUPPET_PACKAGE_NAME)
-    conf_file = util.get_cfg_option_str(
-        puppet_cfg, 'conf_file', PUPPET_CONF_PATH)
-    ssl_dir = util.get_cfg_option_str(puppet_cfg, 'ssl_dir', PUPPET_SSL_DIR)
-    csr_attributes_path = util.get_cfg_option_str(
-        puppet_cfg, 'csr_attributes_path', PUPPET_CSR_ATTRIBUTES_PATH)
+    collection = util.get_cfg_option_str(puppet_cfg, 'collection', None)
+    install_type = util.get_cfg_option_str(
+        puppet_cfg, 'install_type', 'packages')
+    cleanup = util.get_cfg_option_bool(puppet_cfg, 'cleanup', True)
+    run = util.get_cfg_option_bool(puppet_cfg, 'exec', default=False)
+    aio_install_url = util.get_cfg_option_str(
+        puppet_cfg, 'aio_install_url', default=AIO_INSTALL_URL)
 
-    p_constants = PuppetConstants(conf_file, ssl_dir, csr_attributes_path, log)
+    # AIO and distro packages use different paths
+    if install_type == 'aio':
+        puppet_user = 'root'
+        puppet_bin = '/opt/puppetlabs/bin/puppet'
+        puppet_package = 'puppet-agent'
+    else:  # default to 'packages'
+        puppet_user = 'puppet'
+        puppet_bin = 'puppet'
+        puppet_package = 'puppet'
+
+    package_name = util.get_cfg_option_str(
+        puppet_cfg, 'package_name', puppet_package)
     if not install and version:
-        log.warning(("Puppet install set false but version supplied,"
+        log.warning(("Puppet install set to false but version supplied,"
                      " doing nothing."))
     elif install:
-        log.debug(("Attempting to install puppet %s,"),
-                  version if version else 'latest')
+        log.debug(("Attempting to install puppet %s from %s"),
+                  version if version else 'latest', install_type)
 
-        cloud.distro.install_packages((package_name, version))
+        if install_type == "packages":
+            cloud.distro.install_packages((package_name, version))
+        elif install_type == "aio":
+            install_puppet_aio(aio_install_url, version, collection, cleanup)
+        else:
+            log.warning("Unknown puppet install type '%s'", install_type)
+            run = False
+
+    conf_file = util.get_cfg_option_str(
+        puppet_cfg, 'conf_file', get_config_value(puppet_bin, 'config'))
+    ssl_dir = util.get_cfg_option_str(
+        puppet_cfg, 'ssl_dir', get_config_value(puppet_bin, 'ssldir'))
+    csr_attributes_path = util.get_cfg_option_str(
+        puppet_cfg, 'csr_attributes_path',
+        get_config_value(puppet_bin, 'csr_attributes'))
+
+    p_constants = PuppetConstants(conf_file, ssl_dir, csr_attributes_path, log)
 
     # ... and then update the puppet configuration
     if 'conf' in puppet_cfg:
@@ -165,17 +254,18 @@ def handle(name, cfg, cloud, log, _args):
             source=p_constants.conf_path)
         for (cfg_name, cfg) in puppet_cfg['conf'].items():
             # Cert configuration is a special case
-            # Dump the puppet master ca certificate in the correct place
+            # Dump the puppetserver ca certificate in the correct place
             if cfg_name == 'ca_cert':
                 # Puppet ssl sub-directory isn't created yet
                 # Create it with the proper permissions and ownership
                 util.ensure_dir(p_constants.ssl_dir, 0o771)
-                util.chownbyname(p_constants.ssl_dir, 'puppet', 'root')
+                util.chownbyname(p_constants.ssl_dir, puppet_user, 'root')
                 util.ensure_dir(p_constants.ssl_cert_dir)
 
-                util.chownbyname(p_constants.ssl_cert_dir, 'puppet', 'root')
+                util.chownbyname(p_constants.ssl_cert_dir, puppet_user, 'root')
                 util.write_file(p_constants.ssl_cert_path, cfg)
-                util.chownbyname(p_constants.ssl_cert_path, 'puppet', 'root')
+                util.chownbyname(p_constants.ssl_cert_path,
+                                 puppet_user, 'root')
             else:
                 # Iterate through the config items, we'll use ConfigParser.set
                 # to overwrite or create new items as needed
@@ -203,6 +293,25 @@ def handle(name, cfg, cloud, log, _args):
     # Set it up so it autostarts
     _autostart_puppet(log)
 
+    # Run the agent if needed
+    if run:
+        log.debug('Running puppet-agent')
+        cmd = [puppet_bin, 'agent']
+        if 'exec_args' in puppet_cfg:
+            cmd_args = puppet_cfg['exec_args']
+            if isinstance(cmd_args, (list, tuple)):
+                cmd.extend(cmd_args)
+            elif isinstance(cmd_args, str):
+                cmd.extend(cmd_args.split())
+            else:
+                log.warning("Unknown type %s provided for puppet"
+                            " 'exec_args' expected list, tuple,"
+                            " or string", type(cmd_args))
+                cmd.extend(PUPPET_AGENT_DEFAULT_ARGS)
+        else:
+            cmd.extend(PUPPET_AGENT_DEFAULT_ARGS)
+        subp.subp(cmd, capture=False)
+
     # Start puppetd
     subp.subp(['service', 'puppet', 'start'], capture=False)
 
diff --git a/doc/examples/cloud-config-puppet.txt b/doc/examples/cloud-config-puppet.txt
index 3c7e2da7..c6bc15de 100644
--- a/doc/examples/cloud-config-puppet.txt
+++ b/doc/examples/cloud-config-puppet.txt
@@ -1,25 +1,65 @@
 #cloud-config
 #
-# This is an example file to automatically setup and run puppetd
+# This is an example file to automatically setup and run puppet
 # when the instance boots for the first time.
 # Make sure that this file is valid yaml before starting instances.
 # It should be passed as user-data when starting the instance.
 puppet:
+  # Boolean: whether or not to install puppet (default: true)
+  install: true
+
+  # A specific version to pass to the installer script or package manager
+  version: "7.7.0"
+
+  # Valid values are 'packages' and 'aio' (default: 'packages')
+  install_type: "packages"
+
+  # Puppet collection to install if 'install_type' is 'aio'
+  collection: "puppet7"
+
+  # Boolean: whether or not to remove the puppetlabs repo after installation
+  # if 'install_type' is 'aio' (default: true)
+  cleanup: true
+
+  # If 'install_type' is 'aio', change the url to the install script
+  aio_install_url: "https://raw.githubusercontent.com/puppetlabs/install-puppet/main/install.sh"
+
+  # Path to the puppet config file (default: depends on 'install_type')
+  conf_file: "/etc/puppet/puppet.conf"
+
+  # Path to the puppet SSL directory (default: depends on 'install_type')
+  ssl_dir: "/var/lib/puppet/ssl"
+
+  # Path to the CSR attributes file (default: depends on 'install_type')
+  csr_attributes_path: "/etc/puppet/csr_attributes.yaml"
+
+  # The name of the puppet package to install (no-op if 'install_type' is 'aio')
+  package_name: "puppet"
+
+  # Boolean: whether or not to run puppet after configuration finishes
+  # (default: false)
+  exec: false
+
+  # A list of arguments to pass to 'puppet agent' if 'exec' is true
+  # (default: ['--test'])
+  exec_args: ['--test']
+
   # Every key present in the conf object will be added to puppet.conf:
   # [name]
   # subkey=value
   #
   # For example the configuration below will have the following section
   # added to puppet.conf:
-  # [puppetd]
-  # server=puppetmaster.example.org
+  # [main]
+  # server=puppetserver.example.org
   # certname=i-0123456.ip-X-Y-Z.cloud.internal
   #
-  # The puppmaster ca certificate will be available in 
-  # /var/lib/puppet/ssl/certs/ca.pem
+  # The puppetserver ca certificate will be available in
+  # /var/lib/puppet/ssl/certs/ca.pem if using distro packages
+  # or /etc/puppetlabs/puppet/ssl/certs/ca.pem if using AIO packages.
   conf:
     agent:
-      server: "puppetmaster.example.org"
+      server: "puppetserver.example.org"
       # certname supports substitutions at runtime:
       #   %i: instanceid 
       #       Example: i-0123456
@@ -29,11 +69,13 @@ puppet:
       # NB: the certname will automatically be lowercased as required by puppet
       certname: "%i.%f"
     # ca_cert is a special case. It won't be added to puppet.conf.
-    # It holds the puppetmaster certificate in pem format. 
+    # It holds the puppetserver certificate in pem format.
     # It should be a multi-line string (using the | yaml notation for 
     # multi-line strings).
-    # The puppetmaster certificate is located in 
-    # /var/lib/puppet/ssl/ca/ca_crt.pem on the puppetmaster host.
+    # The puppetserver certificate is located in
+    # /var/lib/puppet/ssl/ca/ca_crt.pem on the puppetserver host if using
+    # distro packages or /etc/puppetlabs/puppet/ssl/ca/ca_crt.pem if using AIO
+    # packages.
     #
     ca_cert: |
       -----BEGIN CERTIFICATE-----
diff --git a/tests/cloud_tests/testcases/examples/setup_run_puppet.yaml b/tests/cloud_tests/testcases/examples/setup_run_puppet.yaml
index e366c042..cdb1c28d 100644
--- a/tests/cloud_tests/testcases/examples/setup_run_puppet.yaml
+++ b/tests/cloud_tests/testcases/examples/setup_run_puppet.yaml
@@ -14,14 +14,14 @@ cloud_config: |
    # For example the configuration below will have the following section
    # added to puppet.conf:
    # [puppetd]
-   # server=puppetmaster.example.org
+   # server=puppetserver.example.org
    # certname=i-0123456.ip-X-Y-Z.cloud.internal
    #
    # The puppmaster ca certificate will be available in
    # /var/lib/puppet/ssl/certs/ca.pem
    conf:
      agent:
-       server: "puppetmaster.example.org"
+       server: "puppetserver.example.org"
        # certname supports substitutions at runtime:
        #   %i: instanceid
        #       Example: i-0123456
@@ -31,11 +31,11 @@ cloud_config: |
        # NB: the certname will automatically be lowercased as required by puppet
        certname: "%i.%f"
      # ca_cert is a special case. It won't be added to puppet.conf.
-     # It holds the puppetmaster certificate in pem format.
+     # It holds the puppetserver certificate in pem format.
      # It should be a multi-line string (using the | yaml notation for
      # multi-line strings).
-     # The puppetmaster certificate is located in
-     # /var/lib/puppet/ssl/ca/ca_crt.pem on the puppetmaster host.
+     # The puppetserver certificate is located in
+     # /var/lib/puppet/ssl/ca/ca_crt.pem on the puppetserver host.
      #
      ca_cert: |
        -----BEGIN CERTIFICATE-----
diff --git a/tests/unittests/test_handler/test_handler_puppet.py b/tests/unittests/test_handler/test_handler_puppet.py
index 62388ac6..c0ba2e3c 100644
--- a/tests/unittests/test_handler/test_handler_puppet.py
+++ b/tests/unittests/test_handler/test_handler_puppet.py
@@ -3,8 +3,9 @@
 from cloudinit.config import cc_puppet
 from cloudinit.sources import DataSourceNone
 from cloudinit import (distros, helpers, cloud, util)
-from cloudinit.tests.helpers import CiTestCase, mock
+from cloudinit.tests.helpers import CiTestCase, HttprettyTestCase, mock
 
+import httpretty
 import logging
 import textwrap
 
@@ -63,7 +64,8 @@ class TestPuppetHandle(CiTestCase):
         super(TestPuppetHandle, self).setUp()
         self.new_root = self.tmp_dir()
         self.conf = self.tmp_path('puppet.conf')
-        self.csr_attributes_path = self.tmp_path('csr_attributes.yaml')
+        self.csr_attributes_path = self.tmp_path(
+            'csr_attributes.yaml')
 
     def _get_cloud(self, distro):
         paths = helpers.Paths({'templates_dir': self.new_root})
@@ -72,7 +74,7 @@ class TestPuppetHandle(CiTestCase):
         myds = DataSourceNone.DataSourceNone({}, mydist, paths)
         return cloud.Cloud(myds, paths, {}, mydist, None)
 
-    def test_handler_skips_missing_puppet_key_in_cloudconfig(self, m_auto):
+    def test_skips_missing_puppet_key_in_cloudconfig(self, m_auto):
         """Cloud-config containing no 'puppet' key is skipped."""
         mycloud = self._get_cloud('ubuntu')
         cfg = {}
@@ -81,19 +83,19 @@ class TestPuppetHandle(CiTestCase):
             "no 'puppet' configuration found", self.logs.getvalue())
         self.assertEqual(0, m_auto.call_count)
 
-    @mock.patch('cloudinit.config.cc_puppet.subp.subp')
-    def test_handler_puppet_config_starts_puppet_service(self, m_subp, m_auto):
+    @mock.patch('cloudinit.config.cc_puppet.subp.subp', return_value=("", ""))
+    def test_puppet_config_starts_puppet_service(self, m_subp, m_auto):
         """Cloud-config 'puppet' configuration starts puppet."""
         mycloud = self._get_cloud('ubuntu')
         cfg = {'puppet': {'install': False}}
         cc_puppet.handle('notimportant', cfg, mycloud, LOG, None)
         self.assertEqual(1, m_auto.call_count)
-        self.assertEqual(
+        self.assertIn(
             [mock.call(['service', 'puppet', 'start'], capture=False)],
             m_subp.call_args_list)
 
-    @mock.patch('cloudinit.config.cc_puppet.subp.subp')
-    def test_handler_empty_puppet_config_installs_puppet(self, m_subp, m_auto):
+    @mock.patch('cloudinit.config.cc_puppet.subp.subp', return_value=("", ""))
+    def test_empty_puppet_config_installs_puppet(self, m_subp, m_auto):
         """Cloud-config empty 'puppet' configuration installs latest puppet."""
         mycloud = self._get_cloud('ubuntu')
         mycloud.distro = mock.MagicMock()
@@ -103,8 +105,8 @@ class TestPuppetHandle(CiTestCase):
             [mock.call(('puppet', None))],
             mycloud.distro.install_packages.call_args_list)
 
-    @mock.patch('cloudinit.config.cc_puppet.subp.subp')
-    def test_handler_puppet_config_installs_puppet_on_true(self, m_subp, _):
+    @mock.patch('cloudinit.config.cc_puppet.subp.subp', return_value=("", ""))
+    def test_puppet_config_installs_puppet_on_true(self, m_subp, _):
         """Cloud-config with 'puppet' key installs when 'install' is True."""
         mycloud = self._get_cloud('ubuntu')
         mycloud.distro = mock.MagicMock()
@@ -114,8 +116,85 @@ class TestPuppetHandle(CiTestCase):
             [mock.call(('puppet', None))],
             mycloud.distro.install_packages.call_args_list)
 
-    @mock.patch('cloudinit.config.cc_puppet.subp.subp')
-    def test_handler_puppet_config_installs_puppet_version(self, m_subp, _):
+    @mock.patch('cloudinit.config.cc_puppet.install_puppet_aio', autospec=True)
+    @mock.patch('cloudinit.config.cc_puppet.subp.subp', return_value=("", ""))
+    def test_puppet_config_installs_puppet_aio(self, m_subp, m_aio, _):
+        """Cloud-config with 'puppet' key installs
+        when 'install_type' is 'aio'."""
+        mycloud = self._get_cloud('ubuntu')
+        mycloud.distro = mock.MagicMock()
+        cfg = {'puppet': {'install': True, 'install_type': 'aio'}}
+        cc_puppet.handle('notimportant', cfg, mycloud, LOG, None)
+        m_aio.assert_called_with(
+            cc_puppet.AIO_INSTALL_URL,
+            None, None, True)
+
+    @mock.patch('cloudinit.config.cc_puppet.install_puppet_aio', autospec=True)
+    @mock.patch('cloudinit.config.cc_puppet.subp.subp', return_value=("", ""))
+    def test_puppet_config_installs_puppet_aio_with_version(self,
+                                                            m_subp, m_aio, _):
+        """Cloud-config with 'puppet' key installs
+        when 'install_type' is 'aio' and 'version' is specified."""
+        mycloud = self._get_cloud('ubuntu')
+        mycloud.distro = mock.MagicMock()
+        cfg = {'puppet': {'install': True,
+                          'version': '6.24.0', 'install_type': 'aio'}}
+        cc_puppet.handle('notimportant', cfg, mycloud, LOG, None)
+        m_aio.assert_called_with(
+            cc_puppet.AIO_INSTALL_URL,
+            '6.24.0', None, True)
+
+    @mock.patch('cloudinit.config.cc_puppet.install_puppet_aio', autospec=True)
+    @mock.patch('cloudinit.config.cc_puppet.subp.subp', return_value=("", ""))
+    def test_puppet_config_installs_puppet_aio_with_collection(self,
+                                                               m_subp,
+                                                               m_aio, _):
+        """Cloud-config with 'puppet' key installs
+        when 'install_type' is 'aio' and 'collection' is specified."""
+        mycloud = self._get_cloud('ubuntu')
+        mycloud.distro = mock.MagicMock()
+        cfg = {'puppet': {'install': True,
+                          'collection': 'puppet6', 'install_type': 'aio'}}
+        cc_puppet.handle('notimportant', cfg, mycloud, LOG, None)
+        m_aio.assert_called_with(
+            cc_puppet.AIO_INSTALL_URL,
+            None, 'puppet6', True)
+
+    @mock.patch('cloudinit.config.cc_puppet.install_puppet_aio', autospec=True)
+    @mock.patch('cloudinit.config.cc_puppet.subp.subp', return_value=("", ""))
+    def test_puppet_config_installs_puppet_aio_with_custom_url(self,
+                                                               m_subp,
+                                                               m_aio, _):
+        """Cloud-config with 'puppet' key installs
+        when 'install_type' is 'aio' and 'aio_install_url' is specified."""
+        mycloud = self._get_cloud('ubuntu')
+        mycloud.distro = mock.MagicMock()
+        cfg = {'puppet':
+               {'install': True,
+                'aio_install_url': 'http://test.url/path/to/script.sh',
+                'install_type': 'aio'}}
+        cc_puppet.handle('notimportant', cfg, mycloud, LOG, None)
+        m_aio.assert_called_with(
+            'http://test.url/path/to/script.sh', None, None, True)
+
+    @mock.patch('cloudinit.config.cc_puppet.install_puppet_aio', autospec=True)
+    @mock.patch('cloudinit.config.cc_puppet.subp.subp', return_value=("", ""))
+    def test_puppet_config_installs_puppet_aio_without_cleanup(self,
+                                                               m_subp,
+                                                               m_aio, _):
+        """Cloud-config with 'puppet' key installs
+        when 'install_type' is 'aio' and no cleanup."""
+        mycloud = self._get_cloud('ubuntu')
+        mycloud.distro = mock.MagicMock()
+        cfg = {'puppet': {'install': True,
+                          'cleanup': False, 'install_type': 'aio'}}
+        cc_puppet.handle('notimportant', cfg, mycloud, LOG, None)
+        m_aio.assert_called_with(
+            cc_puppet.AIO_INSTALL_URL,
+            None, None, False)
+
+    @mock.patch('cloudinit.config.cc_puppet.subp.subp', return_value=("", ""))
+    def test_puppet_config_installs_puppet_version(self, m_subp, _):
         """Cloud-config 'puppet' configuration can specify a version."""
         mycloud = self._get_cloud('ubuntu')
         mycloud.distro = mock.MagicMock()
@@ -125,26 +204,39 @@ class TestPuppetHandle(CiTestCase):
             [mock.call(('puppet', '3.8'))],
             mycloud.distro.install_packages.call_args_list)
 
-    @mock.patch('cloudinit.config.cc_puppet.subp.subp')
-    def test_handler_puppet_config_updates_puppet_conf(self, m_subp, m_auto):
+    @mock.patch('cloudinit.config.cc_puppet.get_config_value')
+    @mock.patch('cloudinit.config.cc_puppet.subp.subp', return_value=("", ""))
+    def test_puppet_config_updates_puppet_conf(self,
+                                               m_subp, m_default, m_auto):
         """When 'conf' is provided update values in PUPPET_CONF_PATH."""
+
+        def _fake_get_config_value(puppet_bin, setting):
+            return self.conf
+
+        m_default.side_effect = _fake_get_config_value
         mycloud = self._get_cloud('ubuntu')
         cfg = {
             'puppet': {
-                'conf': {'agent': {'server': 'puppetmaster.example.org'}}}}
-        util.write_file(self.conf, '[agent]\nserver = origpuppet\nother = 3')
-        puppet_conf_path = 'cloudinit.config.cc_puppet.PUPPET_CONF_PATH'
+                'conf': {'agent': {'server': 'puppetserver.example.org'}}}}
+        util.write_file(
+            self.conf, '[agent]\nserver = origpuppet\nother = 3')
         mycloud.distro = mock.MagicMock()
-        with mock.patch(puppet_conf_path, self.conf):
-            cc_puppet.handle('notimportant', cfg, mycloud, LOG, None)
+        cc_puppet.handle('notimportant', cfg, mycloud, LOG, None)
         content = util.load_file(self.conf)
-        expected = '[agent]\nserver = puppetmaster.example.org\nother = 3\n\n'
+        expected = '[agent]\nserver = puppetserver.example.org\nother = 3\n\n'
         self.assertEqual(expected, content)
 
+    @mock.patch('cloudinit.config.cc_puppet.get_config_value')
     @mock.patch('cloudinit.config.cc_puppet.subp.subp')
-    def test_handler_puppet_writes_csr_attributes_file(self, m_subp, m_auto):
+    def test_puppet_writes_csr_attributes_file(self,
+                                               m_subp, m_default, m_auto):
         """When csr_attributes is provided
             creates file in PUPPET_CSR_ATTRIBUTES_PATH."""
+
+        def _fake_get_config_value(puppet_bin, setting):
+            return self.csr_attributes_path
+
+        m_default.side_effect = _fake_get_config_value
         mycloud = self._get_cloud('ubuntu')
         mycloud.distro = mock.MagicMock()
         cfg = {
@@ -163,10 +255,7 @@ class TestPuppetHandle(CiTestCase):
                 }
             }
         }
-        csr_attributes = 'cloudinit.config.cc_puppet.' \
-                         'PUPPET_CSR_ATTRIBUTES_PATH'
-        with mock.patch(csr_attributes, self.csr_attributes_path):
-            cc_puppet.handle('notimportant', cfg, mycloud, LOG, None)
+        cc_puppet.handle('notimportant', cfg, mycloud, LOG, None)
         content = util.load_file(self.csr_attributes_path)
         expected = textwrap.dedent("""\
             custom_attributes:
@@ -177,3 +266,125 @@ class TestPuppetHandle(CiTestCase):
               pp_uuid: ED803750-E3C7-44F5-BB08-41A04433FE2E
             """)
         self.assertEqual(expected, content)
+
+    @mock.patch('cloudinit.config.cc_puppet.subp.subp', return_value=("", ""))
+    def test_puppet_runs_puppet_if_requested(self, m_subp, m_auto):
+        """Run puppet with default args if 'exec' is set to True."""
+        mycloud = self._get_cloud('ubuntu')
+        cfg = {'puppet': {'exec': True}}
+        cc_puppet.handle('notimportant', cfg, mycloud, LOG, None)
+        self.assertEqual(1, m_auto.call_count)
+        self.assertIn(
+            [mock.call(['puppet', 'agent', '--test'], capture=False)],
+            m_subp.call_args_list)
+
+    @mock.patch('cloudinit.config.cc_puppet.subp.subp', return_value=("", ""))
+    def test_puppet_runs_puppet_with_args_list_if_requested(self,
+                                                            m_subp, m_auto):
+        """Run puppet with 'exec_args' list if 'exec' is set to True."""
+        mycloud = self._get_cloud('ubuntu')
+        cfg = {'puppet': {'exec': True, 'exec_args': [
+            '--onetime', '--detailed-exitcodes']}}
+        cc_puppet.handle('notimportant', cfg, mycloud, LOG, None)
+        self.assertEqual(1, m_auto.call_count)
+        self.assertIn(
+            [mock.call(
+                ['puppet', 'agent', '--onetime', '--detailed-exitcodes'],
+                capture=False)],
+            m_subp.call_args_list)
+
+    @mock.patch('cloudinit.config.cc_puppet.subp.subp', return_value=("", ""))
+    def test_puppet_runs_puppet_with_args_string_if_requested(self,
+                                                              m_subp, m_auto):
+        """Run puppet with 'exec_args' string if 'exec' is set to True."""
+        mycloud = self._get_cloud('ubuntu')
+        cfg = {'puppet': {'exec': True,
+                          'exec_args': '--onetime --detailed-exitcodes'}}
+        cc_puppet.handle('notimportant', cfg, mycloud, LOG, None)
+        self.assertEqual(1, m_auto.call_count)
+        self.assertIn(
+            [mock.call(
+                ['puppet', 'agent', '--onetime', '--detailed-exitcodes'],
+                capture=False)],
+            m_subp.call_args_list)
+
+
+class TestInstallPuppetAio(HttprettyTestCase):
+
+    @mock.patch('cloudinit.config.cc_puppet.subp.subp',
+                return_value=(None, None))
+    def test_install_with_default_arguments(self, m_subp):
+        """Install AIO with no arguments"""
+        response = b'#!/bin/bash\necho "Hi Mom"'
+        httpretty.register_uri(
+            httpretty.GET, cc_puppet.AIO_INSTALL_URL,
+            body=response, status=200)
+
+        cc_puppet.install_puppet_aio()
+
+        self.assertEqual(
+            [mock.call([mock.ANY, '--cleanup'], capture=False)],
+            m_subp.call_args_list)
+
+    @mock.patch('cloudinit.config.cc_puppet.subp.subp',
+                return_value=(None, None))
+    def test_install_with_custom_url(self, m_subp):
+        """Install AIO from custom URL"""
+        response = b'#!/bin/bash\necho "Hi Mom"'
+        url = 'http://custom.url/path/to/script.sh'
+        httpretty.register_uri(
+            httpretty.GET, url, body=response, status=200)
+
+        cc_puppet.install_puppet_aio('http://custom.url/path/to/script.sh')
+
+        self.assertEqual(
+            [mock.call([mock.ANY, '--cleanup'], capture=False)],
+            m_subp.call_args_list)
+
+    @mock.patch('cloudinit.config.cc_puppet.subp.subp',
+                return_value=(None, None))
+    def test_install_with_version(self, m_subp):
+        """Install AIO with specific version"""
+        response = b'#!/bin/bash\necho "Hi Mom"'
+        httpretty.register_uri(
+            httpretty.GET, cc_puppet.AIO_INSTALL_URL,
+            body=response, status=200)
+
+        cc_puppet.install_puppet_aio(cc_puppet.AIO_INSTALL_URL, '7.6.0')
+
+        self.assertEqual(
+            [mock.call([mock.ANY, '-v', '7.6.0', '--cleanup'], capture=False)],
+            m_subp.call_args_list)
+
+    @mock.patch('cloudinit.config.cc_puppet.subp.subp',
+                return_value=(None, None))
+    def test_install_with_collection(self, m_subp):
+        """Install AIO with specific collection"""
+        response = b'#!/bin/bash\necho "Hi Mom"'
+        httpretty.register_uri(
+            httpretty.GET, cc_puppet.AIO_INSTALL_URL,
+            body=response, status=200)
+
+        cc_puppet.install_puppet_aio(
+            cc_puppet.AIO_INSTALL_URL, None, 'puppet6-nightly')
+
+        self.assertEqual(
+            [mock.call([mock.ANY, '-c', 'puppet6-nightly', '--cleanup'],
+                       capture=False)],
+            m_subp.call_args_list)
+
+    @mock.patch('cloudinit.config.cc_puppet.subp.subp',
+                return_value=(None, None))
+    def test_install_with_no_cleanup(self, m_subp):
+        """Install AIO with no cleanup"""
+        response = b'#!/bin/bash\necho "Hi Mom"'
+        httpretty.register_uri(
+            httpretty.GET, cc_puppet.AIO_INSTALL_URL,
+            body=response, status=200)
+
+        cc_puppet.install_puppet_aio(
+            cc_puppet.AIO_INSTALL_URL, None, None, False)
+
+        self.assertEqual(
+            [mock.call([mock.ANY], capture=False)],
+            m_subp.call_args_list)