Address review feedback for qube_connection changes

Author: Jayant-kernel

vmupdate/qube_connection.py

@@ -22,14 +22,13 @@
 import shutil
 import signal
 import tempfile
-import asyncio
 import concurrent.futures
 from os.path import join
 from subprocess import CalledProcessError
 from typing import List
 
 import qubesadmin
-from qubesadmin.events.utils import wait_for_domain_shutdown
+import qubesadmin.exc
 from vmupdate.agent.source.args import AgentArgs
 from vmupdate.agent.source.log_config import LOGPATH, LOG_FILE
 from vmupdate.agent.source.status import StatusInfo, FinalStatus
@@ -94,10 +93,15 @@ def __exit__(self, exc_type, exc_val, exc_tb):
                                   self.dest_dir, str(err))
 
         if self.qube.is_running() and not self._initially_running:
-            self.logger.info('Shutdown %s', self.qube.name)
-            self.qube.shutdown()
             if self._has_assigned_pci_devices(self.qube):
-                asyncio.run(wait_for_domain_shutdown([self.qube]))
+                self.logger.info(
+                    'Shutdown %s and wait for full shutdown '
+                    '(PCI devices assigned)', self.qube.name)
+                from vmupdate.vmupdate import shutdown_domains
+                shutdown_domains([self.qube], self.logger)
+            else:
+                self.logger.info('Shutdown %s', self.qube.name)
+                self.qube.shutdown()
 
         self.__connected = False
 
@@ -106,7 +110,7 @@ def _has_assigned_pci_devices(vm) -> bool:
         """Return True when VM has assigned PCI devices."""
         try:
             return any(vm.devices['pci'].get_assigned_devices())
-        except Exception:  # best effort check only
+        except qubesadmin.exc.QubesDaemonAccessError:
             return False
 
     def transfer_agent(self, src_dir: str) -> ProcessResult:

vmupdate/tests/test_qube_connection.py

@@ -2,21 +2,29 @@
 #
 # The Qubes OS Project, https://www.qubes-os.org
 #
-# Copyright (C) 2026
+# Copyright (C) 2025  Jayant Sankar <jayant@example.com>
 #
 # This program is free software; you can redistribute it and/or
 # modify it under the terms of the GNU General Public License
 # as published by the Free Software Foundation; either version 2
 # of the License, or (at your option) any later version.
 #
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program; if not, write to the Free Software
+# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301,
+# USA.
 from unittest.mock import Mock, patch
 
 from vmupdate.qube_connection import QubeConnection
 
 
-@patch("vmupdate.qube_connection.wait_for_domain_shutdown")
-@patch("vmupdate.qube_connection.asyncio.run")
-def test_wait_for_shutdown_when_vm_started_by_update(_arun, wait_shutdown):
+@patch("vmupdate.vmupdate.shutdown_domains")
+def test_wait_for_shutdown_when_vm_started_by_update(shutdown_domains):
     vm = Mock()
     vm.name = "hvm1"
     vm.is_running.side_effect = [False, True]
@@ -30,14 +38,12 @@ def test_wait_for_shutdown_when_vm_started_by_update(_arun, wait_shutdown):
             show_progress=False, status_notifier=status_notifier):
         pass
 
-    vm.shutdown.assert_called_once_with()
-    wait_shutdown.assert_called_once_with([vm])
-    _arun.assert_called_once()
+    shutdown_domains.assert_called_once_with([vm], logger)
+    vm.shutdown.assert_not_called()
 
 
-@patch("vmupdate.qube_connection.wait_for_domain_shutdown")
-@patch("vmupdate.qube_connection.asyncio.run")
-def test_do_not_wait_for_shutdown_without_assigned_pci(_arun, wait_shutdown):
+@patch("vmupdate.vmupdate.shutdown_domains")
+def test_do_not_wait_for_shutdown_without_assigned_pci(shutdown_domains):
     vm = Mock()
     vm.name = "hvm2"
     vm.is_running.side_effect = [False, True]
@@ -52,13 +58,11 @@ def test_do_not_wait_for_shutdown_without_assigned_pci(_arun, wait_shutdown):
         pass
 
     vm.shutdown.assert_called_once_with()
-    wait_shutdown.assert_not_called()
-    _arun.assert_not_called()
+    shutdown_domains.assert_not_called()
 
 
-@patch("vmupdate.qube_connection.wait_for_domain_shutdown")
-@patch("vmupdate.qube_connection.asyncio.run")
-def test_do_not_shutdown_if_vm_was_already_running(_arun, wait_shutdown):
+@patch("vmupdate.vmupdate.shutdown_domains")
+def test_do_not_shutdown_if_vm_was_already_running(shutdown_domains):
     vm = Mock()
     vm.name = "hvm3"
     vm.is_running.return_value = True
@@ -73,5 +77,4 @@ def test_do_not_shutdown_if_vm_was_already_running(_arun, wait_shutdown):
         pass
 
     vm.shutdown.assert_not_called()
-    wait_shutdown.assert_not_called()
-    _arun.assert_not_called()
+    shutdown_domains.assert_not_called()