From 3ecce002975bf364faf2d7c403977f0e1477a44b Mon Sep 17 00:00:00 2001 From: Ben Grande Date: Fri, 22 May 2026 14:51:08 +0200 Subject: [PATCH] Use admin utils wrappers for vm actions --- vmupdate/qube_connection.py | 9 +++--- vmupdate/tests/test_qube_connection.py | 16 ++++------ vmupdate/tests/test_vmupdate.py | 5 ---- vmupdate/utils.py | 39 ++++++++++++++++-------- vmupdate/vmupdate.py | 41 ++++++++++---------------- 5 files changed, 52 insertions(+), 58 deletions(-) diff --git a/vmupdate/qube_connection.py b/vmupdate/qube_connection.py index 9da6b15..9679c0c 100644 --- a/vmupdate/qube_connection.py +++ b/vmupdate/qube_connection.py @@ -18,6 +18,8 @@ # along with this program; if not, write to the Free Software # Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, # USA. + +import asyncio import os import shutil import signal @@ -93,14 +95,13 @@ def __exit__(self, exc_type, exc_val, exc_tb): ) if self.qube.is_running() and not self._initially_running: + wait = False if self._has_assigned_pci_devices(self.qube): self.logger.info( 'Waiting for full shutdown %s (PCI devices assigned)', self.qube.name) - shutdown_domains([self.qube], self.logger) - else: - self.logger.info('Shutdown %s', self.qube.name) - self.qube.shutdown() + wait = True + asyncio.run(shutdown_domains([self.qube], self.logger, wait=wait)) self.__connected = False diff --git a/vmupdate/tests/test_qube_connection.py b/vmupdate/tests/test_qube_connection.py index d8f7669..c0b0a7d 100644 --- a/vmupdate/tests/test_qube_connection.py +++ b/vmupdate/tests/test_qube_connection.py @@ -23,8 +23,7 @@ from vmupdate.qube_connection import QubeConnection -@patch("vmupdate.qube_connection.shutdown_domains") -def test_wait_for_shutdown_when_vm_started_by_update(shutdown_domains): +def test_wait_for_shutdown_when_vm_started_by_update(): vm = Mock() vm.name = "hvm1" vm.is_running.side_effect = [False, True] @@ -38,12 +37,10 @@ def test_wait_for_shutdown_when_vm_started_by_update(shutdown_domains): show_progress=False, status_notifier=status_notifier): pass - shutdown_domains.assert_called_once_with([vm], logger) - vm.shutdown.assert_not_called() + vm.shutdown.assert_called_once_with(force=False, wait=True) -@patch("vmupdate.qube_connection.shutdown_domains") -def test_do_not_wait_for_shutdown_without_assigned_pci(shutdown_domains): +def test_do_not_wait_for_shutdown_without_assigned_pci(): vm = Mock() vm.name = "hvm2" vm.is_running.side_effect = [False, True] @@ -57,12 +54,10 @@ def test_do_not_wait_for_shutdown_without_assigned_pci(shutdown_domains): show_progress=False, status_notifier=status_notifier): pass - vm.shutdown.assert_called_once_with() - shutdown_domains.assert_not_called() + vm.shutdown.assert_called_once_with(force=False, wait=False) -@patch("vmupdate.qube_connection.shutdown_domains") -def test_do_not_shutdown_if_vm_was_already_running(shutdown_domains): +def test_do_not_shutdown_if_vm_was_already_running(): vm = Mock() vm.name = "hvm3" vm.is_running.return_value = True @@ -77,4 +72,3 @@ def test_do_not_shutdown_if_vm_was_already_running(shutdown_domains): pass vm.shutdown.assert_not_called() - shutdown_domains.assert_not_called() diff --git a/vmupdate/tests/test_vmupdate.py b/vmupdate/tests/test_vmupdate.py index 7b0b4d5..b6fcaae 100644 --- a/vmupdate/tests/test_vmupdate.py +++ b/vmupdate/tests/test_vmupdate.py @@ -383,11 +383,9 @@ def test_selection( @patch("vmupdate.update_manager.UpdateAgentManager") @patch("multiprocessing.Pool") @patch("multiprocessing.Manager") -@patch("asyncio.run") @patch("subprocess.Popen") def test_restarting( dummy_subprocess, - arun, mp_manager, mp_pool, agent_mng, @@ -506,7 +504,6 @@ def test_restarting( fails = {args: failed[args] for args in failed if failed[args]} assert not fails - arun.asseert_called() stat = FinalStatus @@ -719,7 +716,6 @@ def test_error( @patch("os.chown") @patch("logging.FileHandler") @patch("logging.getLogger") -@patch("asyncio.run") @pytest.mark.parametrize( "action, code", ( @@ -729,7 +725,6 @@ def test_error( ), ) def test_error_apply( - _arun, _logger, _log_file, _chmod, diff --git a/vmupdate/utils.py b/vmupdate/utils.py index ae76775..23f53bc 100644 --- a/vmupdate/utils.py +++ b/vmupdate/utils.py @@ -22,27 +22,40 @@ from datetime import datetime import qubesadmin.exc -from qubesadmin.events.utils import wait_for_domain_shutdown +from qubesadmin.utils import shutdown, start from vmupdate.agent.source.common.exit_codes import EXIT -def shutdown_domains(to_shutdown, log): +async def shutdown_domains( + to_shutdown, + log, + wait: bool = False, + force: bool = False, +): """ Try to shut down vms and wait to finish. """ ret_code = EXIT.OK - wait_for = [] - for vm in to_shutdown: - try: - vm.shutdown(force=True) - wait_for.append(vm) - except qubesadmin.exc.QubesVMError as exc: - log.error(str(exc)) - ret_code = EXIT.ERR_SHUTDOWN_APP + all_failed = [] + failed = await shutdown(domains=to_shutdown, wait=wait, force=force) + for qube, exc in failed.items(): + log.error(str(exc)) + all_failed.append(qube) + ret_code = EXIT.ERR_SHUTDOWN_APP + done = [qube for qube in to_shutdown if qube not in all_failed] + return ret_code, done - asyncio.run(wait_for_domain_shutdown(wait_for)) - return ret_code, wait_for +async def restart_vms(to_restart, log): + """ + Try to restart vms. + """ + ret_code, shutdowns = await shutdown_domains(to_restart, log) + failed = await start(domains=shutdowns) + for qube, exc in failed.items(): + log.error(str(exc)) + ret_code = EXIT.ERR_START_APP + return ret_code def get_feature(vm, feature_name, default_value=None): @@ -84,3 +97,5 @@ def is_stale(vm, expiration_period): except qubesadmin.exc.QubesDaemonCommunicationError: pass return False + + diff --git a/vmupdate/vmupdate.py b/vmupdate/vmupdate.py index d9fa476..16131a8 100644 --- a/vmupdate/vmupdate.py +++ b/vmupdate/vmupdate.py @@ -4,6 +4,7 @@ """ import argparse +import asyncio import logging import sys import os @@ -11,11 +12,17 @@ from typing import Set, Iterable, Dict, Tuple import qubesadmin +import qubesadmin.utils import qubesadmin.exc from vmupdate.agent.source.status import FinalStatus from vmupdate.agent.source.common.exit_codes import EXIT -from vmupdate.utils import shutdown_domains, get_feature, get_boolean_feature, \ - is_stale +from vmupdate.utils import ( + shutdown_domains, + restart_vms, + get_feature, + get_boolean_feature, + is_stale, +) from . import update_manager from .agent.source.args import AgentArgs @@ -109,9 +116,9 @@ def main(args=None, app=qubesadmin.Qubes()): and no_updates ) - ret_code_restart = apply_updates_to_appvm( + ret_code_restart = asyncio.run(apply_updates_to_appvm( args, independent, templ_statuses, app_statuses, log - ) + )) ret_code = max( ret_code_admin, ret_code_independent, ret_code_appvm, ret_code_restart @@ -386,7 +393,7 @@ def run_update( return ret_code, statuses -def apply_updates_to_appvm( +async def apply_updates_to_appvm( args, vm_updated: Iterable, template_statuses: Dict[str, FinalStatus], @@ -435,7 +442,7 @@ def apply_updates_to_appvm( # first shutdown templates to apply changes to the root volume # they are no need to start templates automatically - ret_code, _ = shutdown_domains(templates_to_shutdown, log) + ret_code, _ = await shutdown_domains(templates_to_shutdown, log) if ret_code != EXIT.OK: log.error("Shutdown of some templates fails with code %d", ret_code) @@ -454,11 +461,11 @@ def apply_updates_to_appvm( ) # both flags `restart` and `apply-to-all` include service vms - ret_code_ = restart_vms(to_restart, log) + ret_code_ = await restart_vms(to_restart, log) ret_code = max(ret_code, ret_code_) if args.apply_to_all: # there is no need to start plain AppVMs automatically - ret_code_, _ = shutdown_domains(to_shutdown, log) + ret_code_, _ = await shutdown_domains(to_shutdown, log) ret_code = max(ret_code, ret_code_) return ret_code @@ -486,23 +493,5 @@ def get_derived_vm_to_apply(templates, derived_statuses): return to_restart, to_shutdown - -def restart_vms(to_restart, log): - """ - Try to restart vms. - """ - ret_code, shutdowns = shutdown_domains(to_restart, log) - - # restart shutdown qubes - for vm in shutdowns: - try: - vm.start() - except qubesadmin.exc.QubesVMError as exc: - log.error(str(exc)) - ret_code = EXIT.ERR_START_APP - - return ret_code - - if __name__ == "__main__": sys.exit(main())