VMware: Support for user provided pre and post-customization scripts

In the VMware customization workflow, we have some options for the user
to upload scripts for additional customization. Based on user request,
those custom scripts can be either run before regular customization or
after. For post customization scripts, we decide whether to run the scripts
just after customization or post system reboot.

4 files changed, 213 insertions, 4 deletions

diff --git a/tests/unittests/test_datasource/test_ovf.py b/tests/unittests/test_datasource/test_ovf.py
index 700da86c..fc4eb36e 100644
--- a/tests/unittests/test_datasource/test_ovf.py
+++ b/tests/unittests/test_datasource/test_ovf.py
@@ -5,11 +5,17 @@
 # This file is part of cloud-init. See LICENSE file for license information.
 
 import base64
-from collections import OrderedDict
+import os
 
-from cloudinit.tests import helpers as test_helpers
+from collections import OrderedDict
+from textwrap import dedent
 
+from cloudinit import util
+from cloudinit.tests.helpers import CiTestCase, wrap_and_call
+from cloudinit.helpers import Paths
 from cloudinit.sources import DataSourceOVF as dsovf
+from cloudinit.sources.helpers.vmware.imc.config_custom_script import (
+    CustomScriptNotFound)
 
 OVF_ENV_CONTENT = """<?xml version="1.0" encoding="UTF-8"?>
 <Environment xmlns="http://schemas.dmtf.org/ovf/environment/1"
@@ -42,7 +48,7 @@ def fill_properties(props, template=OVF_ENV_CONTENT):
     return template.format(properties=properties)
 
 
-class TestReadOvfEnv(test_helpers.TestCase):
+class TestReadOvfEnv(CiTestCase):
     def test_with_b64_userdata(self):
         user_data = "#!/bin/sh\necho hello world\n"
         user_data_b64 = base64.b64encode(user_data.encode()).decode()
@@ -72,7 +78,104 @@ class TestReadOvfEnv(test_helpers.TestCase):
         self.assertIsNone(ud)
 
 
-class TestTransportIso9660(test_helpers.CiTestCase):
+class TestMarkerFiles(CiTestCase):
+
+    def setUp(self):
+        super(TestMarkerFiles, self).setUp()
+        self.tdir = self.tmp_dir()
+
+    def test_false_when_markerid_none(self):
+        """Return False when markerid provided is None."""
+        self.assertFalse(
+            dsovf.check_marker_exists(markerid=None, marker_dir=self.tdir))
+
+    def test_markerid_file_exist(self):
+        """Return False when markerid file path does not exist,
+        True otherwise."""
+        self.assertFalse(
+            dsovf.check_marker_exists('123', self.tdir))
+
+        marker_file = self.tmp_path('.markerfile-123.txt', self.tdir)
+        util.write_file(marker_file, '')
+        self.assertTrue(
+            dsovf.check_marker_exists('123', self.tdir)
+        )
+
+    def test_marker_file_setup(self):
+        """Test creation of marker files."""
+        markerfilepath = self.tmp_path('.markerfile-hi.txt', self.tdir)
+        self.assertFalse(os.path.exists(markerfilepath))
+        dsovf.setup_marker_files(markerid='hi', marker_dir=self.tdir)
+        self.assertTrue(os.path.exists(markerfilepath))
+
+
+class TestDatasourceOVF(CiTestCase):
+
+    with_logs = True
+
+    def setUp(self):
+        super(TestDatasourceOVF, self).setUp()
+        self.datasource = dsovf.DataSourceOVF
+        self.tdir = self.tmp_dir()
+
+    def test_get_data_false_on_none_dmi_data(self):
+        """When dmi for system-product-name is None, get_data returns False."""
+        paths = Paths({'seed_dir': self.tdir})
+        ds = self.datasource(sys_cfg={}, distro={}, paths=paths)
+        retcode = wrap_and_call(
+            'cloudinit.sources.DataSourceOVF',
+            {'util.read_dmi_data': None},
+            ds.get_data)
+        self.assertFalse(retcode, 'Expected False return from ds.get_data')
+        self.assertIn(
+            'DEBUG: No system-product-name found', self.logs.getvalue())
+
+    def test_get_data_no_vmware_customization_disabled(self):
+        """When vmware customization is disabled via sys_cfg log a message."""
+        paths = Paths({'seed_dir': self.tdir})
+        ds = self.datasource(
+            sys_cfg={'disable_vmware_customization': True}, distro={},
+            paths=paths)
+        retcode = wrap_and_call(
+            'cloudinit.sources.DataSourceOVF',
+            {'util.read_dmi_data': 'vmware'},
+            ds.get_data)
+        self.assertFalse(retcode, 'Expected False return from ds.get_data')
+        self.assertIn(
+            'DEBUG: Customization for VMware platform is disabled.',
+            self.logs.getvalue())
+
+    def test_get_data_vmware_customization_disabled(self):
+        """When cloud-init workflow for vmware is enabled via sys_cfg log a
+        message.
+        """
+        paths = Paths({'seed_dir': self.tdir})
+        ds = self.datasource(
+            sys_cfg={'disable_vmware_customization': False}, distro={},
+            paths=paths)
+        conf_file = self.tmp_path('test-cust', self.tdir)
+        conf_content = dedent("""\
+            [CUSTOM-SCRIPT]
+            SCRIPT-NAME = test-script
+            [MISC]
+            MARKER-ID = 12345345
+            """)
+        util.write_file(conf_file, conf_content)
+        with self.assertRaises(CustomScriptNotFound) as context:
+            wrap_and_call(
+                'cloudinit.sources.DataSourceOVF',
+                {'util.read_dmi_data': 'vmware',
+                 'util.del_dir': True,
+                 'search_file': self.tdir,
+                 'wait_for_imc_cfg_file': conf_file,
+                 'get_nics_to_enable': ''},
+                ds.get_data)
+        customscript = self.tmp_path('test-script', self.tdir)
+        self.assertIn('Script %s not found!!' % customscript,
+                      str(context.exception))
+
+
+class TestTransportIso9660(CiTestCase):
 
     def setUp(self):
         super(TestTransportIso9660, self).setUp()
diff --git a/tests/unittests/test_vmware/__init__.py b/tests/unittests/test_vmware/__init__.py
new file mode 100644
index 00000000..e69de29b
--- /dev/null
+++ b/tests/unittests/test_vmware/__init__.py
diff --git a/tests/unittests/test_vmware/test_custom_script.py b/tests/unittests/test_vmware/test_custom_script.py
new file mode 100644
index 00000000..2d9519b0
--- /dev/null
+++ b/tests/unittests/test_vmware/test_custom_script.py
@@ -0,0 +1,99 @@
+# Copyright (C) 2015 Canonical Ltd.
+# Copyright (C) 2017 VMware INC.
+#
+# Author: Maitreyee Saikia <msaikia@vmware.com>
+#
+# This file is part of cloud-init. See LICENSE file for license information.
+
+from cloudinit import util
+from cloudinit.sources.helpers.vmware.imc.config_custom_script import (
+    CustomScriptConstant,
+    CustomScriptNotFound,
+    PreCustomScript,
+    PostCustomScript,
+)
+from cloudinit.tests.helpers import CiTestCase, mock
+
+
+class TestVmwareCustomScript(CiTestCase):
+    def setUp(self):
+        self.tmpDir = self.tmp_dir()
+
+    def test_prepare_custom_script(self):
+        """
+        This test is designed to verify the behavior based on the presence of
+        custom script. Mainly needed for scenario where a custom script is
+        expected, but was not properly copied. "CustomScriptNotFound" exception
+        is raised in such cases.
+        """
+        # Custom script does not exist.
+        preCust = PreCustomScript("random-vmw-test", self.tmpDir)
+        self.assertEqual("random-vmw-test", preCust.scriptname)
+        self.assertEqual(self.tmpDir, preCust.directory)
+        self.assertEqual(self.tmp_path("random-vmw-test", self.tmpDir),
+                         preCust.scriptpath)
+        with self.assertRaises(CustomScriptNotFound):
+            preCust.prepare_script()
+
+        # Custom script exists.
+        custScript = self.tmp_path("test-cust", self.tmpDir)
+        util.write_file(custScript, "test-CR-strip/r/r")
+        postCust = PostCustomScript("test-cust", self.tmpDir)
+        self.assertEqual("test-cust", postCust.scriptname)
+        self.assertEqual(self.tmpDir, postCust.directory)
+        self.assertEqual(custScript, postCust.scriptpath)
+        self.assertFalse(postCust.postreboot)
+        postCust.prepare_script()
+        # Check if all carraige returns are stripped from script.
+        self.assertFalse("/r" in custScript)
+
+    def test_rc_local_exists(self):
+        """
+        This test is designed to verify the different scenarios associated
+        with the presence of rclocal.
+        """
+        # test when rc local does not exist
+        postCust = PostCustomScript("test-cust", self.tmpDir)
+        with mock.patch.object(CustomScriptConstant, "RC_LOCAL", "/no/path"):
+            rclocal = postCust.find_rc_local()
+            self.assertEqual("", rclocal)
+
+        # test when rc local exists
+        rclocalFile = self.tmp_path("vmware-rclocal", self.tmpDir)
+        util.write_file(rclocalFile, "# Run post-reboot guest customization",
+                        omode="w")
+        with mock.patch.object(CustomScriptConstant, "RC_LOCAL", rclocalFile):
+            rclocal = postCust.find_rc_local()
+            self.assertEqual(rclocalFile, rclocal)
+            self.assertTrue(postCust.has_previous_agent, rclocal)
+
+        # test when rc local is a symlink
+        rclocalLink = self.tmp_path("dummy-rclocal-link", self.tmpDir)
+        util.sym_link(rclocalFile, rclocalLink, True)
+        with mock.patch.object(CustomScriptConstant, "RC_LOCAL", rclocalLink):
+            rclocal = postCust.find_rc_local()
+            self.assertEqual(rclocalFile, rclocal)
+
+    def test_execute_post_cust(self):
+        """
+        This test is to identify if rclocal was properly populated to be
+        run after reboot.
+        """
+        customscript = self.tmp_path("vmware-post-cust-script", self.tmpDir)
+        rclocal = self.tmp_path("vmware-rclocal", self.tmpDir)
+        # Create a temporary rclocal file
+        open(customscript, "w")
+        util.write_file(rclocal, "tests\nexit 0", omode="w")
+        postCust = PostCustomScript("vmware-post-cust-script", self.tmpDir)
+        with mock.patch.object(CustomScriptConstant, "RC_LOCAL", rclocal):
+            # Test that guest customization agent is not installed initially.
+            self.assertFalse(postCust.postreboot)
+            self.assertIs(postCust.has_previous_agent(rclocal), False)
+            postCust.install_agent()
+
+            # Assert rclocal has been modified to have guest customization
+            # agent.
+            self.assertTrue(postCust.postreboot)
+            self.assertTrue(postCust.has_previous_agent, rclocal)
+
+# vi: ts=4 expandtab
diff --git a/tests/unittests/test_vmware_config_file.py b/tests/unittests/test_vmware_config_file.py
index 0f8cda95..036f6879 100644
--- a/tests/unittests/test_vmware_config_file.py
+++ b/tests/unittests/test_vmware_config_file.py
@@ -335,5 +335,12 @@ class TestVmwareConfigFile(CiTestCase):
         self.assertEqual('255.255.0.0', subnet.get('netmask'),
                          'Subnet netmask')
 
+    def test_custom_script(self):
+        cf = ConfigFile("tests/data/vmware/cust-dhcp-2nic.cfg")
+        conf = Config(cf)
+        self.assertIsNone(conf.custom_script_name)
+        cf._insertKey("CUSTOM-SCRIPT|SCRIPT-NAME", "test-script")
+        conf = Config(cf)
+        self.assertEqual("test-script", conf.custom_script_name)
 
 # vi: ts=4 expandtab