From bb1db7d7fd99d6f0f76ec2987196ab8ec33262d4 Mon Sep 17 00:00:00 2001
From: Robert Jordens <jordens@gmail.com>
Date: Thu, 28 Jan 2016 16:54:14 -0700
Subject: [PATCH] test.ctlmgr: add basic test tooling

---
 artiq/test/ctlmgr.py | 70 ++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 70 insertions(+)
 create mode 100644 artiq/test/ctlmgr.py

diff --git a/artiq/test/ctlmgr.py b/artiq/test/ctlmgr.py
new file mode 100644
index 000000000..f45f8ace7
--- /dev/null
+++ b/artiq/test/ctlmgr.py
@@ -0,0 +1,70 @@
+import os
+import unittest
+import logging
+import asyncio
+
+from artiq.devices.ctlmgr import Controllers
+from artiq.protocols.pc_rpc import AsyncioClient
+
+logger = logging.getLogger(__name__)
+
+
+class ControllerCase(unittest.TestCase):
+    def setUp(self):
+        if os.name == "nt":
+            self.loop = asyncio.ProactorEventLoop()
+        else:
+            self.loop = asyncio.new_event_loop()
+        asyncio.set_event_loop(self.loop)
+        self.addCleanup(self.loop.close)
+
+        self.controllers = Controllers()
+        self.controllers.host_filter = "::1"
+        self.addCleanup(
+            lambda: self.loop.run_until_complete(self.controllers.shutdown()))
+
+    async def start(self, name, entry):
+        self.controllers[name] = entry
+        await self.controllers.queue.join()
+        await self.wait_for_ping(entry["host"], entry["port"])
+
+    async def get_client(self, host, port):
+        remote = AsyncioClient()
+        await remote.connect_rpc(host, port, None)
+        targets, _ = remote.get_rpc_id()
+        remote.select_rpc_target(targets[0])
+        self.addCleanup(remote.close_rpc)
+        return remote
+
+    async def wait_for_ping(self, host, port, retries=5, timeout=2):
+        dt = timeout/retries
+        while timeout > 0:
+            try:
+                remote = await self.get_client(host, port)
+                ok = await asyncio.wait_for(remote.ping(), dt)
+                if not ok:
+                    raise ValueError("unexcepted ping() response from "
+                                     "controller: `{}`".format(ok))
+                return ok
+            except asyncio.TimeoutError:
+                timeout -= dt
+            except (ConnectionAbortedError, ConnectionError,
+                    ConnectionRefusedError, ConnectionResetError):
+                await asyncio.sleep(dt)
+                timeout -= dt
+        raise asyncio.TimeoutError
+
+    def test_start_ping_stop_controller(self):
+        entry = {
+            "type": "controller",
+            "host": "::1",
+            "port": 3253,
+            "command": "lda_controller -p {port} --bind {bind} "
+                    "--no-localhost-bind --simulation",
+        }
+        async def test():
+            await self.start("lda_sim", entry)
+            remote = await self.get_client(entry["host"], entry["port"])
+            await remote.close()
+
+        self.loop.run_until_complete(test())