tools: move event loop creation/closing to main function

Do not close event loop in utility function - handle it only in main().
For this reason, change appropriate functions to coroutines.

Fixes QubesOS/qubes-issues#2865

qubesadmin/events/utils.py

@@ -42,19 +42,20 @@ def interrupt_on_vm_shutdown(vm, subject, event):
         raise Interrupt
 
 
-def wait_for_domain_shutdown(vm, timeout):
+@asyncio.coroutine
+def wait_for_domain_shutdown(vm, timeout, loop=None):
     ''' Helper function to wait for domain shutdown.
 
     This function wait for domain shutdown, but do not initiate the shutdown
     itself.
 
-    Note: you need to close event loop after calling this function.
-
     :param vm: QubesVM object to wait for shutdown on
     :param timeout: Timeout in seconds, use 0 for no timeout
+    :param loop: asyncio event loop
     '''
-    events = qubesadmin.events.EventsDispatcher(vm.app)
+    if loop is None:
         loop = asyncio.get_event_loop()
+    events = qubesadmin.events.EventsDispatcher(vm.app)
     events.add_handler('domain-shutdown',
         functools.partial(interrupt_on_vm_shutdown, vm))
     events.add_handler('connection-established',
@@ -65,7 +66,7 @@ def wait_for_domain_shutdown(vm, timeout):
         # pylint: disable=no-member
         loop.call_later(timeout, events_task.cancel)
     try:
-        loop.run_until_complete(events_task)
+        yield from events_task
     except asyncio.CancelledError:
         raise qubesadmin.exc.QubesVMShutdownTimeout(
             'VM %s shutdown timeout expired', vm.name)

qubesadmin/tests/tools/qvm_template_postprocess.py

@@ -255,6 +255,7 @@ class TC_00_qvm_template_postprocess(qubesadmin.tests.QubesTestCase):
                 ('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],
             app=self.app)
@@ -274,6 +275,10 @@ class TC_00_qvm_template_postprocess(qubesadmin.tests.QubesTestCase):
         ])
         self.assertAllCalled()
 
+    @asyncio.coroutine
+    def wait_for_shutdown(self, vm, timeout):
+        pass
+
     @mock.patch('qubesadmin.tools.qvm_template_postprocess.import_appmenus')
     @mock.patch('qubesadmin.tools.qvm_template_postprocess.import_root_img')
     def test_021_post_install_reinstall(self, mock_import_root_img,
@@ -298,11 +303,13 @@ class TC_00_qvm_template_postprocess(qubesadmin.tests.QubesTestCase):
                 '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],
             app=self.app)
@@ -334,7 +341,9 @@ class TC_00_qvm_template_postprocess(qubesadmin.tests.QubesTestCase):
                 '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',
             self.source_dir.name],

qubesadmin/tools/qvm_template_postprocess.py

@@ -143,6 +143,7 @@ def import_appmenus(vm, source_dir):
         vm.log.warning('Failed to set default application list: %s', e)
 
 
+@asyncio.coroutine
 def post_install(args):
     '''Handle post-installation tasks'''
 
@@ -192,17 +193,16 @@ def post_install(args):
             if have_events:
                 try:
                     # pylint: disable=no-member
-                    qubesadmin.events.utils.wait_for_domain_shutdown(vm,
+                    yield from qubesadmin.events.utils.wait_for_domain_shutdown(
-                        qubesadmin.config.defaults['shutdown_timeout'])
+                        vm, qubesadmin.config.defaults['shutdown_timeout'])
                 except qubesadmin.exc.QubesVMShutdownTimeout:
                     vm.kill()
-                asyncio.get_event_loop().close()
             else:
                 timeout = qubesadmin.config.defaults['shutdown_timeout']
                 while timeout >= 0:
                     if vm.is_halted():
                         break
-                    time.sleep(1)
+                    yield from asyncio.sleep(1)
                     timeout -= 1
                 if not vm.is_halted():
                     vm.kill()
@@ -252,7 +252,13 @@ def main(args=None, app=None):
     if not args.really:
         parser.error('Do not call this tool directly.')
     if args.action == 'post-install':
-        return post_install(args)
+        loop = asyncio.get_event_loop()
+        try:
+            loop.run_until_complete(post_install(args))
+            loop.stop()
+            loop.run_forever()
+        finally:
+            loop.close()
     elif args.action == 'pre-remove':
         pre_remove(args)
     else: