From bf78e0c7d242f6b9f159451eb444b1608f678c2d Mon Sep 17 00:00:00 2001
From: Sebastien Bourdeauducq <sb@m-labs.hk>
Date: Thu, 9 Aug 2018 16:51:12 +0800
Subject: [PATCH] test: fix handling of missing devices

---
 artiq/test/coredevice/test_moninj.py | 15 ++++++++++++---
 artiq/test/hardware_testbench.py     | 10 +++++-----
 2 files changed, 17 insertions(+), 8 deletions(-)

diff --git a/artiq/test/coredevice/test_moninj.py b/artiq/test/coredevice/test_moninj.py
index bda590bc8..3d9481ae0 100644
--- a/artiq/test/coredevice/test_moninj.py
+++ b/artiq/test/coredevice/test_moninj.py
@@ -1,3 +1,4 @@
+import unittest
 import asyncio
 
 from artiq.coredevice.comm_moninj import *
@@ -6,9 +7,17 @@ from artiq.test.hardware_testbench import ExperimentCase
 
 class MonInjTest(ExperimentCase):
     def test_moninj(self):
-        core_host = self.device_mgr.get_desc("core")["arguments"]["host"]
-        loop_out_channel = self.device_mgr.get_desc("loop_out")["arguments"]["channel"]
-        loop_in_channel = self.device_mgr.get_desc("loop_in")["arguments"]["channel"]
+        try:
+            core = self.device_mgr.get_desc("core")
+            loop_out = self.device_mgr.get_desc("loop_out")
+            loop_in = self.device_mgr.get_desc("loop_in")
+        except KeyError as e:
+            # skip if ddb does not match requirements
+            raise unittest.SkipTest(
+                "test device not available: `{}`".format(*e.args))
+        core_host = core["arguments"]["host"]
+        loop_out_channel = loop_out["arguments"]["channel"]
+        loop_in_channel = loop_in["arguments"]["channel"]
 
         notifications = []
         injection_statuses = []
diff --git a/artiq/test/hardware_testbench.py b/artiq/test/hardware_testbench.py
index 4ce3114c3..32bc8cd4a 100644
--- a/artiq/test/hardware_testbench.py
+++ b/artiq/test/hardware_testbench.py
@@ -11,7 +11,7 @@ import time
 import socket
 
 from artiq.master.databases import DeviceDB, DatasetDB
-from artiq.master.worker_db import DeviceManager, DatasetManager
+from artiq.master.worker_db import DeviceManager, DatasetManager, DeviceError
 from artiq.coredevice.core import CompileError
 from artiq.frontend.artiq_run import DummyScheduler
 from artiq.protocols.pc_rpc import AutoTarget, Client
@@ -113,12 +113,12 @@ class ExperimentCase(unittest.TestCase):
             exp = cls(
                 (self.device_mgr, self.dataset_mgr, None),
                 *args, **kwargs)
-            exp.prepare()
-            return exp
-        except KeyError as e:
+        except DeviceError as e:
             # skip if ddb does not match requirements
             raise unittest.SkipTest(
-                "device_db entry `{}` not found".format(*e.args))
+                "test device not available: `{}`".format(*e.args))
+        exp.prepare()
+        return exp
 
     def execute(self, cls, *args, **kwargs):
         expid = {