Extract _install_dism_drivers helper

Signed-off-by: Mihaela Balutoiu <mbalutoiu@cloudbasesolutions.com>

coriolis/osmorphing/windows.py

@@ -810,27 +810,37 @@ class BaseWindowsMorphingTools(base.BaseOSMorphingTools):
                     virtio_drive, d, virtio_dir, arch) for d in drivers
             ]
 
-            sid = self._get_sid()
-            # Fails on Nano Server without explicitly granting permissions
-            file_repo_path = (
-                "%sWindows\\System32\\DriverStore\\FileRepository" %
-                self._os_root_dir
-            )
-            self._grant_permissions(file_repo_path, "*%s" % sid)
-            try:
-                for driver_path in driver_paths:
-                    if self._conn.test_path(driver_path):
-                        self._add_dism_driver(driver_path)
-                    else:
-                        LOG.warn(
-                            "Could not locate driver dir '%s', skipping.",
-                            driver_path
-                        )
-            finally:
-                self._revoke_permissions(file_repo_path, "*%s" % sid)
+            self._install_dism_drivers(driver_paths)
         finally:
             self._dismount_disk_image(VIRTIO_WIN_ISO_PATH)
 
+    def _install_dism_drivers(self, driver_paths):
+        """Installs the drivers located at the given paths into the offline
+        Windows image via DISM.
+
+        Temporarily grants permissions on the DriverStore FileRepository
+        (required on Nano Server) for the duration of the installation.
+        Paths which do not exist are skipped with a warning.
+        """
+        sid = self._get_sid()
+        # Fails on Nano Server without explicitly granting permissions
+        file_repo_path = (
+            "%sWindows\\System32\\DriverStore\\FileRepository" %
+            self._os_root_dir
+        )
+        self._grant_permissions(file_repo_path, "*%s" % sid)
+        try:
+            for driver_path in driver_paths:
+                if self._conn.test_path(driver_path):
+                    self._add_dism_driver(driver_path)
+                else:
+                    LOG.warning(
+                        "Could not locate driver dir '%s', skipping.",
+                        driver_path
+                    )
+        finally:
+            self._revoke_permissions(file_repo_path, "*%s" % sid)
+
     def get_packages(self):
         return [], []
 

coriolis/tests/osmorphing/test_windows.py

@@ -1227,3 +1227,51 @@ class BaseWindowsMorphingToolsTestCase(test_base.CoriolisBaseTestCase):
 
         self.morphing_tools._conn.download_file.assert_not_called()
         mock_mount_disk_image.assert_not_called()
+
+    @mock.patch.object(windows.BaseWindowsMorphingTools, "_get_sid")
+    @mock.patch.object(windows.BaseWindowsMorphingTools, "_grant_permissions")
+    @mock.patch.object(windows.BaseWindowsMorphingTools, "_add_dism_driver")
+    @mock.patch.object(windows.BaseWindowsMorphingTools, "_revoke_permissions")
+    def test_install_dism_drivers(
+            self, mock_revoke_permissions, mock_add_dism_driver,
+            mock_grant_permissions, mock_get_sid):
+        null_sid = "S-1-0-0"
+        mock_get_sid.return_value = null_sid
+
+        existing_driver_path = "e:\\NetKVM\\w11\\amd64"
+        missing_driver_path = "e:\\Missing\\w11\\amd64"
+        driver_paths = [existing_driver_path, missing_driver_path]
+        self.conn.test_path.side_effect = lambda path: (
+            path == existing_driver_path)
+
+        self.morphing_tools._install_dism_drivers(driver_paths)
+
+        exp_repo_path = "C:\\Windows\\System32\\DriverStore\\FileRepository"
+        mock_grant_permissions.assert_called_once_with(
+            exp_repo_path, f"*{null_sid}")
+        mock_revoke_permissions.assert_called_once_with(
+            exp_repo_path, f"*{null_sid}")
+        # Only the existing driver path should be installed:
+        mock_add_dism_driver.assert_called_once_with(existing_driver_path)
+
+    @mock.patch.object(windows.BaseWindowsMorphingTools, "_get_sid")
+    @mock.patch.object(windows.BaseWindowsMorphingTools, "_grant_permissions")
+    @mock.patch.object(windows.BaseWindowsMorphingTools, "_add_dism_driver")
+    @mock.patch.object(windows.BaseWindowsMorphingTools, "_revoke_permissions")
+    def test_install_dism_drivers_revokes_on_failure(
+            self, mock_revoke_permissions, mock_add_dism_driver,
+            mock_grant_permissions, mock_get_sid):
+        null_sid = "S-1-0-0"
+        mock_get_sid.return_value = null_sid
+        mock_add_dism_driver.side_effect = IOError
+
+        self.assertRaises(
+            IOError,
+            self.morphing_tools._install_dism_drivers, ["e:\\NetKVM"])
+
+        exp_repo_path = "C:\\Windows\\System32\\DriverStore\\FileRepository"
+        mock_grant_permissions.assert_called_once_with(
+            exp_repo_path, f"*{null_sid}")
+        # Permissions must be revoked even if driver installation fails:
+        mock_revoke_permissions.assert_called_once_with(
+            exp_repo_path, f"*{null_sid}")