vmupdate: wait for VM shutdown before leaving update connection

Author: Jayant-kernel

vmupdate/qube_connection.py

@@ -22,12 +22,14 @@
 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
 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,6 +96,7 @@ def __exit__(self, exc_type, exc_val, exc_tb):
         if self.qube.is_running() and not self._initially_running:
             self.logger.info('Shutdown %s', self.qube.name)
             self.qube.shutdown()
+            asyncio.run(wait_for_domain_shutdown([self.qube]))
 
         self.__connected = False
 

vmupdate/tests/test_qube_connection.py

@@ -0,0 +1,52 @@
+# coding=utf-8
+#
+# The Qubes OS Project, https://www.qubes-os.org
+#
+# Copyright (C) 2026
+#
+# 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.
+#
+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):
+    vm = Mock()
+    vm.name = "hvm1"
+    vm.is_running.side_effect = [False, True]
+    status_notifier = Mock()
+    logger = Mock()
+
+    with QubeConnection(
+            vm, "/tmp/qubes-update", cleanup=False, logger=logger,
+            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()
+
+
+@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):
+    vm = Mock()
+    vm.name = "hvm2"
+    vm.is_running.return_value = True
+    status_notifier = Mock()
+    logger = Mock()
+
+    with QubeConnection(
+            vm, "/tmp/qubes-update", cleanup=False, logger=logger,
+            show_progress=False, status_notifier=status_notifier):
+        pass
+
+    vm.shutdown.assert_not_called()
+    wait_shutdown.assert_not_called()
+    _arun.assert_not_called()