1 files changed, 107 insertions, 1 deletions
diff --git a/tests/pa/test_deprovision.py b/tests/pa/test_deprovision.py
index be34915..c4cd9b4 100644
--- a/tests/pa/test_deprovision.py
+++ b/tests/pa/test_deprovision.py
@@ -15,11 +15,117 @@
 # Requires Python 2.4+ and Openssl 1.0+
 #
 
-from tests.tools import *
+import tempfile
+
+import azurelinuxagent.common.utils.fileutil as fileutil
+
 from azurelinuxagent.pa.deprovision import get_deprovision_handler
+from azurelinuxagent.pa.deprovision.default import DeprovisionHandler
+from tests.tools import *
 
 
 class TestDeprovision(AgentTestCase):
+    @patch("azurelinuxagent.pa.deprovision.default.DeprovisionHandler.cloud_init_dirs")
+    @patch("azurelinuxagent.pa.deprovision.default.DeprovisionHandler.cloud_init_files")
+    def test_del_cloud_init_without_once(self,
+                mock_files,
+                mock_dirs):
+        deprovision_handler = get_deprovision_handler("","","")
+        deprovision_handler.del_cloud_init([], [], include_once=False)
+
+        mock_dirs.assert_called_with(include_once=False)
+        mock_files.assert_called_with(include_once=False)
+
+    @patch("signal.signal")
+    @patch("azurelinuxagent.common.protocol.get_protocol_util")
+    @patch("azurelinuxagent.common.osutil.get_osutil")
+    @patch("azurelinuxagent.pa.deprovision.default.DeprovisionHandler.cloud_init_dirs")
+    @patch("azurelinuxagent.pa.deprovision.default.DeprovisionHandler.cloud_init_files")
+    def test_del_cloud_init(self,
+                mock_files,
+                mock_dirs,
+                mock_osutil,
+                mock_util,
+                mock_signal):
+        try:
+            with tempfile.NamedTemporaryFile() as f:
+                warnings = []
+                actions = []
+
+                dirs = [tempfile.mkdtemp()]
+                mock_dirs.return_value = dirs
+
+                files = [f.name]
+                mock_files.return_value = files
+
+                deprovision_handler = get_deprovision_handler("","","")
+                deprovision_handler.del_cloud_init(warnings, actions)
+
+                mock_dirs.assert_called_with(include_once=True)
+                mock_files.assert_called_with(include_once=True)
+
+                self.assertEqual(len(warnings), 0)
+                self.assertEqual(len(actions), 2)
+                for da in actions:
+                    if da.func == fileutil.rm_dirs:
+                        self.assertEqual(da.args, dirs)
+                    elif da.func == fileutil.rm_files:
+                        self.assertEqual(da.args, files)
+                    else:
+                        self.assertTrue(False)
+
+                try:
+                    for da in actions:
+                        da.invoke()
+                    self.assertEqual(len([d for d in dirs if os.path.isdir(d)]), 0)
+                    self.assertEqual(len([f for f in files if os.path.isfile(f)]), 0)
+                except Exception as e:
+                    self.assertTrue(False, "Exception {0}".format(e))
+        except OSError:
+            # Ignore the error caused by removing the file within the "with"
+            pass
+    
+    @distros("ubuntu")
+    @patch('azurelinuxagent.common.conf.get_lib_dir')
+    def test_del_lib_dir_files(self,
+                        distro_name,
+                        distro_version,
+                        distro_full_name,
+                        mock_conf):
+        files = [
+            'HostingEnvironmentConfig.xml',
+            'Incarnation',
+            'Protocol',
+            'SharedConfig.xml',
+            'WireServerEndpoint',
+            'Extensions.1.xml',
+            'ExtensionsConfig.1.xml',
+            'GoalState.1.xml',
+            'Extensions.2.xml',
+            'ExtensionsConfig.2.xml',
+            'GoalState.2.xml'
+        ]
+
+        tmp = tempfile.mkdtemp()
+        mock_conf.return_value = tmp
+        for f in files:
+            fileutil.write_file(os.path.join(tmp, f), "Value")
+
+        deprovision_handler = get_deprovision_handler(distro_name,
+                                                      distro_version,
+                                                      distro_full_name)
+        warnings = []
+        actions = []
+        deprovision_handler.del_lib_dir_files(warnings, actions)
+
+        self.assertTrue(len(warnings) == 0)
+        self.assertTrue(len(actions) == 1)
+        self.assertEqual(fileutil.rm_files, actions[0].func)
+        self.assertTrue(len(actions[0].args) > 0)
+        for f in actions[0].args:
+            self.assertTrue(os.path.basename(f) in files)
+
+
     @distros("redhat")
     def test_deprovision(self,
                          distro_name,