From 2bb4ad189a37e71b53aa353c4386b5ca9e820b67 Mon Sep 17 00:00:00 2001
From: Sebastien Bourdeauducq <sb@m-labs.hk>
Date: Tue, 22 Mar 2016 22:29:41 +0800
Subject: [PATCH] test: fix controller simulations

---
 artiq/test/hardware_testbench.py  | 14 ++++++++---
 artiq/test/test_novatech409b.py   | 22 +++++++++++++---
 artiq/test/test_thorlabs_tcube.py | 42 +++++++++++++++++++++++++------
 3 files changed, 64 insertions(+), 14 deletions(-)

diff --git a/artiq/test/hardware_testbench.py b/artiq/test/hardware_testbench.py
index 890f7e0e8..a36dea032 100644
--- a/artiq/test/hardware_testbench.py
+++ b/artiq/test/hardware_testbench.py
@@ -21,10 +21,12 @@ artiq_root = os.getenv("ARTIQ_ROOT")
 logger = logging.getLogger(__name__)
 
 
-@unittest.skipUnless(artiq_root, "no ARTIQ_ROOT")
-class ControllerCase(unittest.TestCase):
+class GenericControllerCase(unittest.TestCase):
+    def get_device_db(self):
+        raise NotImplementedError
+
     def setUp(self):
-        self.device_db = DeviceDB(os.path.join(artiq_root, "device_db.pyon"))
+        self.device_db = self.get_device_db()
         self.device_mgr = DeviceManager(self.device_db)
         self.controllers = {}
 
@@ -87,6 +89,12 @@ class ControllerCase(unittest.TestCase):
             del self.controllers[name]
 
 
+@unittest.skipUnless(artiq_root, "no ARTIQ_ROOT")
+class ControllerCase(GenericControllerCase):
+    def get_device_db(self):
+        return DeviceDB(os.path.join(artiq_root, "device_db.pyon"))
+
+
 @unittest.skipUnless(artiq_root, "no ARTIQ_ROOT")
 class ExperimentCase(unittest.TestCase):
     def setUp(self):
diff --git a/artiq/test/test_novatech409b.py b/artiq/test/test_novatech409b.py
index 8c6327ce4..2a18890eb 100644
--- a/artiq/test/test_novatech409b.py
+++ b/artiq/test/test_novatech409b.py
@@ -1,7 +1,7 @@
+import sys
 import unittest
 
-from artiq.devices.novatech409b.driver import Novatech409B
-from artiq.test.hardware_testbench import ControllerCase
+from artiq.test.hardware_testbench import GenericControllerCase, ControllerCase
 
 
 class GenericNovatech409BTest:
@@ -27,6 +27,20 @@ class TestNovatech409B(GenericNovatech409BTest, ControllerCase):
         self.driver = self.device_mgr.get("novatech409b")
 
 
-class TestNovatech409BSim(GenericNovatech409BTest, unittest.TestCase):
+class TestNovatech409BSim(GenericNovatech409BTest, GenericControllerCase):
+    def get_device_db(self):
+        return {
+            "novatech409b": {
+                "type": "controller",
+                "host": "::1",
+                "port": 3254,
+                "command": (sys.executable.replace("\\", "\\\\")
+                            + " -m artiq.frontend.novatech409b_controller "
+                            + "-p {port} --simulation")
+            }
+        }
+
     def setUp(self):
-        self.driver = Novatech409B(None)
+        GenericControllerCase.setUp(self)
+        self.start_controller("novatech409b")
+        self.driver = self.device_mgr.get("novatech409b")
diff --git a/artiq/test/test_thorlabs_tcube.py b/artiq/test/test_thorlabs_tcube.py
index bf04a7a06..f435da544 100644
--- a/artiq/test/test_thorlabs_tcube.py
+++ b/artiq/test/test_thorlabs_tcube.py
@@ -1,9 +1,9 @@
-import unittest
 import time
+import sys
+import unittest
 
-from artiq.devices.thorlabs_tcube.driver import TdcSim, TpzSim
 from artiq.language.units import V
-from artiq.test.hardware_testbench import ControllerCase
+from artiq.test.hardware_testbench import ControllerCase, GenericControllerCase
 
 
 class GenericTdcTest:
@@ -138,9 +138,23 @@ class TestTdc(ControllerCase, GenericTdcTest):
         self.cont = self.device_mgr.get("tdc")
 
 
-class TestTdcSim(unittest.TestCase, GenericTdcTest):
+class TestTdcSim(GenericControllerCase, GenericTdcTest):
+    def get_device_db(self):
+        return {
+            "tdc": {
+                "type": "controller",
+                "host": "::1",
+                "port": 3255,
+                "command": (sys.executable.replace("\\", "\\\\")
+                            + " -m artiq.frontend.thorlabs_tcube_controller "
+                            + "-p {port} -P tdc001 --simulation")
+            }
+        }
+
     def setUp(self):
-        self.cont = TdcSim()
+        GenericControllerCase.setUp(self)
+        self.start_controller("tdc")
+        self.cont = self.device_mgr.get("tdc")
 
 
 class TestTpz(ControllerCase, GenericTpzTest):
@@ -150,6 +164,20 @@ class TestTpz(ControllerCase, GenericTpzTest):
         self.cont = self.device_mgr.get("tpz")
 
 
-class TestTpzSim(unittest.TestCase, GenericTpzTest):
+class TestTpzSim(GenericControllerCase, GenericTpzTest):
+    def get_device_db(self):
+        return {
+            "tpz": {
+                "type": "controller",
+                "host": "::1",
+                "port": 3255,
+                "command": (sys.executable.replace("\\", "\\\\")
+                            + " -m artiq.frontend.thorlabs_tcube_controller "
+                            + "-p {port} -P tpz001 --simulation")
+            }
+        }
+
     def setUp(self):
-        self.cont = TpzSim()
+        GenericControllerCase.setUp(self)
+        self.start_controller("tpz")
+        self.cont = self.device_mgr.get("tpz")