Make qvm_shutdown reusable for other tools

It was necessary to make alternative main async functions, as Python is
picky about running async code when mixing async and sync functions.

Author: ben-grande

qubesadmin/tests/tools/qvm_kill.py

@@ -20,6 +20,8 @@
 
 # pylint: disable=missing-docstring
 
+import asyncio
+
 import qubesadmin.tests
 import qubesadmin.tests.tools
 import qubesadmin.tools.qvm_kill
@@ -35,6 +37,20 @@ def test_000_with_vm(self):
         qubesadmin.tools.qvm_kill.main(['some-vm'], app=self.app)
         self.assertAllCalled()
 
+    def test_000_with_vm_async(self):
+        loop = asyncio.new_event_loop()
+        asyncio.set_event_loop(loop)
+
+        self.app.expected_calls[
+            ('dom0', 'admin.vm.List', None, None)] = \
+            b'0\x00some-vm class=AppVM state=Running\n'
+        self.app.expected_calls[
+            ('some-vm', 'admin.vm.Kill', None, None)] = b'0\x00'
+        loop.run_until_complete(
+            qubesadmin.tools.qvm_kill.main_async(['some-vm'], app=self.app)
+        )
+        self.assertAllCalled()
+
     def test_001_missing_vm(self):
         with self.assertRaises(SystemExit):
             with qubesadmin.tests.tools.StderrBuffer() as stderr:
@@ -74,9 +90,9 @@ def test_004_other_error(self):
             ('dom0', 'admin.vm.List', None, None)] = \
             b'0\x00some-vm class=AppVM state=Running\n'
         with qubesadmin.tests.tools.StderrBuffer() as stderr:
-            self.assertEqual(
-                qubesadmin.tools.qvm_kill.main(['some-vm'], app=self.app),
-                1)
+            with self.assertRaises(SystemExit):
+                self.assertEqual(
+                    qubesadmin.tools.qvm_kill.main(['some-vm'], app=self.app),
+                    1)
         self.assertAllCalled()
-        self.assertIn("Failed to kill 'some-vm': Error message",
-                      stderr.getvalue())
+        self.assertIn("Failed to kill: some-vm", stderr.getvalue())

qubesadmin/tests/tools/qvm_shutdown.py

@@ -96,6 +96,24 @@ def test_010_wait(self):
         qubesadmin.tools.qvm_shutdown.main(['--wait', 'some-vm'], app=self.app)
         self.assertAllCalled()
 
+    def test_011_wait_async(self):
+        '''test --wait option with async main'''
+        loop = asyncio.new_event_loop()
+        asyncio.set_event_loop(loop)
+
+        self.app.expected_calls[
+            ('some-vm', 'admin.vm.Shutdown', 'wait', None)] = \
+            b'0\x00'
+        self.app.expected_calls[
+            ('dom0', 'admin.vm.List', None, None)] = \
+            b'0\x00some-vm class=AppVM state=Running\n'
+        loop.run_until_complete(
+            qubesadmin.tools.qvm_shutdown.main_async(
+                ['--wait', 'some-vm'], app=self.app
+            )
+        )
+        self.assertAllCalled()
+
     def test_012_wait_all(self):
         '''test --wait option, with multiple VMs'''
         loop = asyncio.new_event_loop()
@@ -198,6 +216,21 @@ def test_006_dry_run(self):
             ['--dry-run', 'some-vm'], app=self.app)
         self.assertAllCalled()
 
+    def test_006_dry_run_async(self):
+        '''test --dry-run skips shutdown calls'''
+        loop = asyncio.new_event_loop()
+        asyncio.set_event_loop(loop)
+
+        self.app.expected_calls[
+            ('dom0', 'admin.vm.List', None, None)] = \
+            b'0\x00some-vm class=AppVM state=Running\n'
+        loop.run_until_complete(
+            qubesadmin.tools.qvm_shutdown.main_async(
+                ['--dry-run', 'some-vm'], app=self.app
+            )
+        )
+        self.assertAllCalled()
+
     def test_011_wait_retry(self):
         '''test --wait retries VMs whose shutdown request failed'''
         loop = asyncio.new_event_loop()

qubesadmin/tests/tools/qvm_template_postprocess.py

@@ -402,17 +402,6 @@ def test_020_post_install(self, mock_import_root_img,
         self.app.expected_calls[
             ('test-vm', 'admin.vm.Shutdown', None, None)] = b'0\0'
 
-        if qubesadmin.tools.qvm_template_postprocess.have_events:
-            patch_domain_shutdown = mock.patch(
-                'qubesadmin.events.utils.wait_for_domain_shutdown')
-            self.addCleanup(patch_domain_shutdown.stop)
-            mock_domain_shutdown = patch_domain_shutdown.start()
-            mock_domain_shutdown.side_effect = self.wait_for_shutdown
-        else:
-            self.app.expected_calls[
-                ('test-vm', 'admin.vm.List', None, None)] = \
-                b'0\0test-vm class=TemplateVM state=Halted\n'
-
         asyncio.set_event_loop(asyncio.new_event_loop())
         ret = qubesadmin.tools.qvm_template_postprocess.main([
             '--really', 'post-install', 'test-vm', self.source_dir.name],
@@ -424,9 +413,6 @@ def test_020_post_install(self, mock_import_root_img,
             'test-vm'], self.source_dir.name)
         mock_import_appmenus.assert_called_once_with(self.app.domains[
             'test-vm'], self.source_dir.name, skip_generate=True)
-        if qubesadmin.tools.qvm_template_postprocess.have_events:
-            mock_domain_shutdown.assert_called_once_with([self.app.domains[
-                'test-vm']])
         self.assertEqual(self.app.service_calls, [
             ('test-vm', 'qubes.PostInstall', {
                 'stdin': subprocess.PIPE,
@@ -460,17 +446,6 @@ def test_021_post_install_reinstall(self, mock_reset_private_img,
         self.app.expected_calls[
             ('test-vm', 'admin.vm.Shutdown', None, None)] = b'0\0'
 
-        if qubesadmin.tools.qvm_template_postprocess.have_events:
-            patch_domain_shutdown = mock.patch(
-                'qubesadmin.events.utils.wait_for_domain_shutdown')
-            self.addCleanup(patch_domain_shutdown.stop)
-            mock_domain_shutdown = patch_domain_shutdown.start()
-            mock_domain_shutdown.side_effect = self.wait_for_shutdown
-        else:
-            self.app.expected_calls[
-                ('test-vm', 'admin.vm.List', None, None)] = \
-                b'0\0test-vm class=TemplateVM state=Halted\n'
-
         asyncio.set_event_loop(asyncio.new_event_loop())
         ret = qubesadmin.tools.qvm_template_postprocess.main([
             '--really', 'post-install', 'test-vm', self.source_dir.name],
@@ -483,9 +458,6 @@ def test_021_post_install_reinstall(self, mock_reset_private_img,
             'test-vm'])
         mock_import_appmenus.assert_called_once_with(self.app.domains[
             'test-vm'], self.source_dir.name, skip_generate=True)
-        if qubesadmin.tools.qvm_template_postprocess.have_events:
-            mock_domain_shutdown.assert_called_once_with([self.app.domains[
-                'test-vm']])
         self.assertEqual(self.app.service_calls, [
             ('test-vm', 'qubes.PostInstall', {
                 'stdin': subprocess.PIPE,
@@ -508,13 +480,6 @@ def test_022_post_install_skip_start(self, mock_reset_private_img,
             = b'0\0'
         self.app.add_new_vm = mock.Mock()
 
-        if qubesadmin.tools.qvm_template_postprocess.have_events:
-            patch_domain_shutdown = mock.patch(
-                'qubesadmin.events.utils.wait_for_domain_shutdown')
-            self.addCleanup(patch_domain_shutdown.stop)
-            mock_domain_shutdown = patch_domain_shutdown.start()
-            mock_domain_shutdown.side_effect = self.wait_for_shutdown
-
         asyncio.set_event_loop(asyncio.new_event_loop())
         ret = qubesadmin.tools.qvm_template_postprocess.main([
             '--really', '--skip-start', 'post-install', 'test-vm',
@@ -528,8 +493,6 @@ def test_022_post_install_skip_start(self, mock_reset_private_img,
             'test-vm'])
         mock_import_appmenus.assert_called_once_with(self.app.domains[
             'test-vm'], self.source_dir.name, skip_generate=False)
-        if qubesadmin.tools.qvm_template_postprocess.have_events:
-            self.assertFalse(mock_domain_shutdown.called)
         self.assertEqual(self.app.service_calls, [])
         self.assertAllCalled()
 

qubesadmin/tools/qvm_kill.py

@@ -21,36 +21,44 @@
 '''Immediately terminate a qube without a graceful shutdown sequence.'''
 
 
+import asyncio
 import sys
+
 import qubesadmin.exc
 import qubesadmin.tools
+import qubesadmin.tools.qvm_shutdown
 
 parser = qubesadmin.tools.QubesArgumentParser(
     description='immediately terminate a qube without a graceful shutdown'
                 ' sequence',
     vmname_nargs='+')
 
-def main(args=None, app=None):
-    '''Main routine of :program:`qvm-kill`.
 
-    :param list args: Optional arguments to override those delivered from \
-        command line.
-    '''
+async def run_async(domains=None):
+    # pylint: disable=missing-docstring
+    remnants = await qubesadmin.tools.qvm_shutdown.kill(domains=domains)
+    if not remnants:
+        return 0
+    parser.error_runtime(
+        "Failed to kill: {}".format(
+            ", ".join(qube.name for qube in remnants)
+        ),
+        len(remnants)
+    )
+
 
+async def main_async(args=None, app=None):
+    # pylint: disable=missing-docstring
     args = parser.parse_args(args, app=app)
+    return await run_async(domains=args.domains)
 
-    exit_code = 0
-    for domain in args.domains:
-        try:
-            domain.kill()
-        except qubesadmin.exc.QubesVMNotStartedError:
-            pass
-        except (IOError, OSError, qubesadmin.exc.QubesException) as e:
-            exit_code = 1
-            parser.print_error("Failed to kill '{}': {}".format(
-                domain.name, e))
-
-    return exit_code
+
+def main(args=None, app=None):
+    # pylint: disable=missing-docstring
+    args = parser.parse_args(args, app=app)
+    loop = asyncio.new_event_loop()
+    asyncio.set_event_loop(loop)
+    return loop.run_until_complete(run_async(domains=args.domains))
 
 
 if __name__ == '__main__':

qubesadmin/tools/qvm_shutdown.py

@@ -70,12 +70,7 @@
 )
 
 
-def shutdown(
-    args,
-    loop: asyncio.AbstractEventLoop,
-    domains: list[qubesadmin.vm.QubesVM],
-    force: bool,
-):
+async def shutdown(args, domains: list[qubesadmin.vm.QubesVM]):
     """
     Asynchronously shutdown qubes and return qubes that failed to shutdown as
     well as failed with a timeout.
@@ -84,14 +79,12 @@ def shutdown(
     remnants, timedout = [], []
     tasks = [
         asyncio.wait_for(
-            async_thread(qube.shutdown, force=force, wait=args.wait),
+            async_thread(qube.shutdown, force=args.force, wait=args.wait),
             timeout=args.timeout,
         )
         for qube in domains
     ]
-    results = loop.run_until_complete(
-        asyncio.gather(*tasks, return_exceptions=True)
-    )
+    results = await asyncio.gather(*tasks, return_exceptions=True)
     for qube, res in zip(domains, results):
         if not isinstance(res, BaseException):
             qube.log.info("Shutdown succeeded")
@@ -120,16 +113,14 @@ def shutdown(
     return remnants, timedout
 
 
-def kill(loop: asyncio.AbstractEventLoop, domains: list[qubesadmin.vm.QubesVM]):
+async def kill(domains: list[qubesadmin.vm.QubesVM]):
     """
     Asynchronously kill qubes and return qubes that failed to shutdown.
     """
     # pylint: disable=missing-docstring
     remnants = domains.copy()
     tasks = [async_thread(qube.kill) for qube in domains]
-    results = loop.run_until_complete(
-        asyncio.gather(*tasks, return_exceptions=True)
-    )
+    results = await asyncio.gather(*tasks, return_exceptions=True)
     for qube, res in zip(domains, results):
         if not isinstance(res, BaseException):
             qube.log.info("Killing succeeded")
@@ -144,36 +135,24 @@ def kill(loop: asyncio.AbstractEventLoop, domains: list[qubesadmin.vm.QubesVM]):
     return remnants
 
 
-def main(args=None, app=None):
+async def run_async(args=None, domains=None):
     # pylint: disable=missing-docstring
-    args = parser.parse_args(args, app=app)
-    force = args.force or (args.all_domains and not args.exclude)
-    if args.dry_run:
-        return
-    domains = set(args.domains)
-
-    loop = asyncio.new_event_loop()
-    asyncio.set_event_loop(loop)
-    remnants, timedout = shutdown(
-        args=args, force=force, loop=loop, domains=domains
-    )
+    remnants, timedout = await shutdown(args=args, domains=domains)
     if timedout:
         args.app.log.info(
             "Retrying shutdown of qubes that timed out: {}".format(
                 ", ".join(qube.name for qube in timedout)
             )
         )
-        remnants, timedout = shutdown(
-            args=args, force=force, loop=loop, domains=timedout
-        )
+        remnants, timedout = await shutdown(args=args, domains=timedout)
 
     if timedout:
         args.app.log.info(
             "Killing timed out qubes: {}".format(
                 ", ".join(qube.name for qube in timedout)
             )
         )
-        remnants = kill(loop=loop, domains=timedout)
+        remnants = await kill(domains=timedout)
 
     if not remnants:
         return
@@ -185,5 +164,31 @@ def main(args=None, app=None):
     )
 
 
+def parse_common(args=None, app=None):
+    # pylint: disable=missing-docstring
+    args = parser.parse_args(args, app=app)
+    args.force = args.force or (args.all_domains and not args.exclude)
+    domains = set(args.domains)
+    return args, domains
+
+
+async def main_async(args=None, app=None):
+    # pylint: disable=missing-docstring
+    args, domains = parse_common(args=args, app=app)
+    if args.dry_run:
+        return
+    await run_async(args=args, domains=domains)
+
+
+def main(args=None, app=None):
+    # pylint: disable=missing-docstring
+    args, domains = parse_common(args=args, app=app)
+    if args.dry_run:
+        return
+    loop = asyncio.new_event_loop()
+    asyncio.set_event_loop(loop)
+    loop.run_until_complete(run_async(args=args, domains=domains))
+
+
 if __name__ == "__main__":
     sys.exit(main())