From 60bdf74137dff57562595403a0da150822e8bb29 Mon Sep 17 00:00:00 2001
From: Yann Sionneau <ys@m-labs.hk>
Date: Thu, 4 Jun 2015 13:40:13 +0200
Subject: [PATCH] tests: use try/finally to close event loop + wait for process
 to die after killing it

---
 artiq/master/worker.py |  4 ++++
 artiq/test/pc_rpc.py   |  6 ++++--
 artiq/test/worker.py   | 35 ++++++++++++++++++++++-------------
 3 files changed, 30 insertions(+), 15 deletions(-)

diff --git a/artiq/master/worker.py b/artiq/master/worker.py
index bb89907f4..30890d485 100644
--- a/artiq/master/worker.py
+++ b/artiq/master/worker.py
@@ -101,6 +101,8 @@ class Worker:
                 logger.warning("failed to send terminate command to worker"
                                " (RID %d), killing", self.rid, exc_info=True)
                 self.process.kill()
+                # Wait for process to terminate
+                yield from self.process.communicate()
                 return
             try:
                 yield from asyncio_process_wait_timeout(self.process,
@@ -108,6 +110,8 @@ class Worker:
             except asyncio.TimeoutError:
                 logger.warning("worker did not exit (RID %d), killing", self.rid)
                 self.process.kill()
+                # Wait for process to terminate
+                yield from self.process.communicate()
             else:
                 logger.debug("worker exited gracefully (RID %d)", self.rid)
         finally:
diff --git a/artiq/test/pc_rpc.py b/artiq/test/pc_rpc.py
index c0afcf6f9..b9e76e19a 100644
--- a/artiq/test/pc_rpc.py
+++ b/artiq/test/pc_rpc.py
@@ -75,8 +75,10 @@ class RPCCase(unittest.TestCase):
     def _loop_asyncio_echo(self):
         loop = asyncio.new_event_loop()
         asyncio.set_event_loop(loop)
-        loop.run_until_complete(self._asyncio_echo())
-        loop.close()
+        try:
+            loop.run_until_complete(self._asyncio_echo())
+        finally:
+            loop.close()
 
     def test_asyncio_echo(self):
         self._run_server_and_test(self._loop_asyncio_echo)
diff --git a/artiq/test/worker.py b/artiq/test/worker.py
index 2bc3a30e3..a72547b8b 100644
--- a/artiq/test/worker.py
+++ b/artiq/test/worker.py
@@ -40,23 +40,29 @@ def _call_worker(worker, expid):
 
 
 def _run_experiment(experiment):
-    loop = asyncio.new_event_loop()
-    asyncio.set_event_loop(loop)
-    expid = {
-        "file": sys.modules[__name__].__file__,
-        "experiment": experiment,
-        "arguments": dict()
-    }
-    handlers = {
-        "init_rt_results": lambda description: None
-    }
+    try:
+        expid = {
+            "file": sys.modules[__name__].__file__,
+            "experiment": experiment,
+            "arguments": dict()
+        }
+        handlers = {
+            "init_rt_results": lambda description: None
+        }
 
-    worker = Worker(handlers)
-    loop.run_until_complete(_call_worker(worker, expid))
-    loop.close()
+        loop = asyncio.get_event_loop()
+        worker = Worker(handlers)
+        loop.run_until_complete(_call_worker(worker, expid))
+    finally:
+        loop.close()
 
 
 class WatchdogCase(unittest.TestCase):
+
+    def setUp(self):
+        self.loop = asyncio.new_event_loop()
+        asyncio.set_event_loop(self.loop)
+
     def test_watchdog_no_timeout(self):
         _run_experiment("WatchdogNoTimeout")
 
@@ -67,3 +73,6 @@ class WatchdogCase(unittest.TestCase):
     def test_watchdog_timeout_in_build(self):
         with self.assertRaises(WorkerWatchdogTimeout):
             _run_experiment("WatchdogTimeoutInBuild")
+
+    def tearDown(self):
+        self.loop.close()