Merge pull request #1642 from vrdmr/vameru/remove-zips-after-uninstall

Cleaning zip packages correctly during uninstall/update.

azurelinuxagent/ga/exthandlers.py

@@ -82,6 +82,7 @@ class ExtCommandEnvVariable(object):
     ExtensionVersion = "AZURE_GUEST_AGENT_EXTENSION_VERSION"
     ExtensionSeqNumber = "ConfigSequenceNumber"  # At par with Windows Guest Agent
 
+
 def get_traceback(e):
     if sys.version_info[0] == 3:
         return e.__traceback__
@@ -849,7 +850,7 @@ def download(self):
         if self.pkg is None or self.pkg.uris is None or len(self.pkg.uris) == 0:
             raise ExtensionDownloadError("No package uri found")
 
-        destination = os.path.join(conf.get_lib_dir(), os.path.basename(self.pkg.uris[0].uri) + ".zip")
+        destination = os.path.join(conf.get_lib_dir(), self.get_extension_package_zipfile_name())
 
         package_exists = False
         if os.path.exists(destination):
@@ -989,16 +990,14 @@ def uninstall(self):
 
     def remove_ext_handler(self):
         try:
-            zip_filename = "__".join(os.path.basename(self.get_base_dir()).split("-")) + ".zip"
-            destination = os.path.join(conf.get_lib_dir(), zip_filename)
-            if os.path.exists(destination):
-                self.pkg_file = destination
-                os.remove(self.pkg_file)
+            zip_filename = os.path.join(conf.get_lib_dir(), self.get_extension_package_zipfile_name())
+            if os.path.exists(zip_filename):
+                os.remove(zip_filename)
+                self.logger.verbose("Deleted the extension zip at path {0}", zip_filename)
 
             base_dir = self.get_base_dir()
             if os.path.isdir(base_dir):
-                self.logger.info("Remove extension handler directory: {0}",
-                                 base_dir)
+                self.logger.info("Remove extension handler directory: {0}", base_dir)
 
                 # some extensions uninstall asynchronously so ignore error 2 while removing them
                 def on_rmtree_error(_, __, exc_info):
@@ -1366,6 +1365,11 @@ def get_handler_status(self):
         except (IOError, ValueError) as e:
             self.logger.error("Failed to get handler status: {0}", e)
 
+    def get_extension_package_zipfile_name(self):
+        return "{0}__{1}{2}".format(self.ext_handler.name,
+                                    self.ext_handler.properties.version,
+                                    HANDLER_PKG_EXT)
+
     def get_full_name(self):
         return "{0}-{1}".format(self.ext_handler.name,
                                 self.ext_handler.properties.version)

tests/ga/test_extension.py

@@ -321,6 +321,14 @@ def _assert_handler_status(self, report_vm_status, expected_status,
         self.assertEquals(expected_ext_count, len(handler_status.extensions))
         return
 
+    def _assert_ext_pkg_file_status(self, expected_to_be_present=True, extension_version="1.0.0",
+                                    extension_handler_name="OSTCExtensions.ExampleHandlerLinux"):
+        zip_file_format = "{0}__{1}.zip"
+        if expected_to_be_present:
+            self.assertIn(zip_file_format.format(extension_handler_name, extension_version), os.listdir(conf.get_lib_dir()))
+        else:
+            self.assertNotIn(zip_file_format.format(extension_handler_name, extension_version), os.listdir(conf.get_lib_dir()))
+
     def _assert_no_handler_status(self, report_vm_status):
         self.assertTrue(report_vm_status.called)
         args, kw = report_vm_status.call_args
@@ -435,6 +443,114 @@ def test_ext_handler(self, *args):
         exthandlers_handler.run()
         self._assert_no_handler_status(protocol.report_vm_status)
 
+    def test_ext_zip_file_packages_removed_in_update_case(self, *args):
+        test_data = WireProtocolData(DATA_FILE)
+        exthandlers_handler, protocol = self._create_mock(test_data, *args)
+
+        # Test enable scenario.
+        exthandlers_handler.run()
+        self._assert_handler_status(protocol.report_vm_status, "Ready", 1, "1.0.0")
+        self._assert_ext_status(protocol.report_ext_status, "success", 0)
+        self._assert_ext_pkg_file_status(expected_to_be_present=True,
+                                         extension_version="1.0.0")
+
+        # Update the package
+        test_data.goal_state = test_data.goal_state.replace("<Incarnation>1<",
+                                                            "<Incarnation>2<")
+        test_data.ext_conf = test_data.ext_conf.replace("seqNo=\"0\"",
+                                                        "seqNo=\"1\"")
+        test_data.ext_conf = test_data.ext_conf.replace("1.0.0", "1.1.0")
+        exthandlers_handler.run()
+        self._assert_handler_status(protocol.report_vm_status, "Ready", 1, "1.1.0")
+        self._assert_ext_status(protocol.report_ext_status, "success", 1)
+        self._assert_ext_pkg_file_status(expected_to_be_present=False,
+                                         extension_version="1.0.0")
+        self._assert_ext_pkg_file_status(expected_to_be_present=True,
+                                         extension_version="1.1.0")
+
+        # Update the package second time
+        test_data.goal_state = test_data.goal_state.replace("<Incarnation>2<",
+                                                            "<Incarnation>3<")
+        test_data.ext_conf = test_data.ext_conf.replace("seqNo=\"1\"",
+                                                        "seqNo=\"2\"")
+        test_data.ext_conf = test_data.ext_conf.replace("1.1.0", "1.2.0")
+        exthandlers_handler.run()
+        self._assert_handler_status(protocol.report_vm_status, "Ready", 1, "1.2.0")
+        self._assert_ext_status(protocol.report_ext_status, "success", 2)
+        self._assert_ext_pkg_file_status(expected_to_be_present=False,
+                                         extension_version="1.1.0")
+        self._assert_ext_pkg_file_status(expected_to_be_present=True,
+                                         extension_version="1.2.0")
+
+    def test_ext_zip_file_packages_removed_in_uninstall_case(self, *args):
+        test_data = WireProtocolData(DATA_FILE)
+        exthandlers_handler, protocol = self._create_mock(test_data, *args)
+        extension_version = "1.0.0"
+
+        # Test enable scenario.
+        exthandlers_handler.run()
+        self._assert_handler_status(protocol.report_vm_status, "Ready", 1, extension_version)
+        self._assert_ext_status(protocol.report_ext_status, "success", 0)
+        self._assert_ext_pkg_file_status(expected_to_be_present=True,
+                                         extension_version=extension_version)
+
+        # Test uninstall
+        test_data.goal_state = test_data.goal_state.replace("<Incarnation>1<",
+                                                            "<Incarnation>2<")
+        test_data.ext_conf = test_data.ext_conf.replace("enabled", "uninstall")
+        exthandlers_handler.run()
+        self._assert_no_handler_status(protocol.report_vm_status)
+        self._assert_ext_pkg_file_status(expected_to_be_present=False,
+                                         extension_version=extension_version)
+
+    def test_ext_zip_file_packages_removed_in_update_and_uninstall_case(self, *args):
+        test_data = WireProtocolData(DATA_FILE)
+        exthandlers_handler, protocol = self._create_mock(test_data, *args)
+
+        # Test enable scenario.
+        exthandlers_handler.run()
+        self._assert_handler_status(protocol.report_vm_status, "Ready", 1, "1.0.0")
+        self._assert_ext_status(protocol.report_ext_status, "success", 0)
+        self._assert_ext_pkg_file_status(expected_to_be_present=True,
+                                         extension_version="1.0.0")
+
+        # Update the package
+        test_data.goal_state = test_data.goal_state.replace("<Incarnation>1<",
+                                                            "<Incarnation>2<")
+        test_data.ext_conf = test_data.ext_conf.replace("seqNo=\"0\"",
+                                                        "seqNo=\"1\"")
+        test_data.ext_conf = test_data.ext_conf.replace("1.0.0", "1.1.0")
+        exthandlers_handler.run()
+        self._assert_handler_status(protocol.report_vm_status, "Ready", 1, "1.1.0")
+        self._assert_ext_status(protocol.report_ext_status, "success", 1)
+        self._assert_ext_pkg_file_status(expected_to_be_present=False,
+                                         extension_version="1.0.0")
+        self._assert_ext_pkg_file_status(expected_to_be_present=True,
+                                         extension_version="1.1.0")
+
+        # Update the package second time
+        test_data.goal_state = test_data.goal_state.replace("<Incarnation>2<",
+                                                            "<Incarnation>3<")
+        test_data.ext_conf = test_data.ext_conf.replace("seqNo=\"1\"",
+                                                        "seqNo=\"2\"")
+        test_data.ext_conf = test_data.ext_conf.replace("1.1.0", "1.2.0")
+        exthandlers_handler.run()
+        self._assert_handler_status(protocol.report_vm_status, "Ready", 1, "1.2.0")
+        self._assert_ext_status(protocol.report_ext_status, "success", 2)
+        self._assert_ext_pkg_file_status(expected_to_be_present=False,
+                                         extension_version="1.1.0")
+        self._assert_ext_pkg_file_status(expected_to_be_present=True,
+                                         extension_version="1.2.0")
+
+        # Test uninstall
+        test_data.goal_state = test_data.goal_state.replace("<Incarnation>3<",
+                                                            "<Incarnation>4<")
+        test_data.ext_conf = test_data.ext_conf.replace("enabled", "uninstall")
+        exthandlers_handler.run()
+        self._assert_no_handler_status(protocol.report_vm_status)
+        self._assert_ext_pkg_file_status(expected_to_be_present=False,
+                                         extension_version="1.2.0")
+
     def test_ext_handler_no_settings(self, *args):
         test_data = WireProtocolData(DATA_FILE_EXT_NO_SETTINGS)
         exthandlers_handler, protocol = self._create_mock(test_data, *args)

tests/ga/test_exthandlers_download_extension.py

@@ -86,7 +86,7 @@ def _create_invalid_zip_file(filename):
             file.write("An invalid ZIP file\n")
 
     def _get_extension_package_file(self):
-        return os.path.join(self.agent_dir, os.path.basename(self.pkg.uris[0].uri) + ".zip")
+        return os.path.join(self.agent_dir, self.ext_handler_instance.get_extension_package_zipfile_name())
 
     def _get_extension_command_file(self):
         return os.path.join(self.extension_dir, DownloadExtensionTestCase._extension_command)

tests/ga/test_exthandlers_exthandlerinstance.py

@@ -2,7 +2,8 @@
 # Licensed under the Apache License.
 
 from azurelinuxagent.ga.exthandlers import ExtHandlerInstance
-from azurelinuxagent.common.protocol.restapi import ExtHandler, ExtHandlerProperties
+from azurelinuxagent.common.protocol.restapi import ExtHandler, ExtHandlerProperties, ExtHandlerPackage, \
+    ExtHandlerVersionUri
 from tests.tools import *
 
 
@@ -16,13 +17,31 @@ def setUp(self):
         ext_handler.properties = ext_handler_properties
         self.ext_handler_instance = ExtHandlerInstance(ext_handler=ext_handler, protocol=None)
 
+        pkg_uri = ExtHandlerVersionUri()
+        pkg_uri.uri = "http://bar/foo__1.2.3"
+        self.ext_handler_instance.pkg = ExtHandlerPackage(ext_handler_properties.version)
+        self.ext_handler_instance.pkg.uris.append(pkg_uri)
+
+        self.base_dir = self.tmp_dir
         self.extension_directory = os.path.join(self.tmp_dir, "extension_directory")
         self.mock_get_base_dir = patch.object(self.ext_handler_instance, "get_base_dir", return_value=self.extension_directory)
         self.mock_get_base_dir.start()
 
     def tearDown(self):
         self.mock_get_base_dir.stop()
 
+    def test_rm_ext_handler_dir_should_remove_the_extension_packages(self):
+        os.mkdir(self.extension_directory)
+        open(os.path.join(self.extension_directory, "extension_file1"), 'w').close()
+        open(os.path.join(self.extension_directory, "extension_file2"), 'w').close()
+        open(os.path.join(self.extension_directory, "extension_file3"), 'w').close()
+        open(os.path.join(self.base_dir, "foo__1.2.3.zip"), 'w').close()
+
+        self.ext_handler_instance.remove_ext_handler()
+
+        self.assertFalse(os.path.exists(self.extension_directory))
+        self.assertFalse(os.path.exists(os.path.join(self.base_dir, "foo__1.2.3.zip")))
+
     def test_rm_ext_handler_dir_should_remove_the_extension_directory(self):
         os.mkdir(self.extension_directory)
         os.mknod(os.path.join(self.extension_directory, "extension_file1"))