diff --git a/artiq/test/pc_rpc.py b/artiq/test/pc_rpc.py index 19000c659..423b4f31d 100644 --- a/artiq/test/pc_rpc.py +++ b/artiq/test/pc_rpc.py @@ -17,12 +17,12 @@ test_object = [5, 2.1, None, True, False, class RPCCase(unittest.TestCase): - def _run_server_and_test(self, test): + def _run_server_and_test(self, test, *args): # running this file outside of unittest starts the echo server with subprocess.Popen([sys.executable, sys.modules[__name__].__file__]) as proc: try: - test() + test(*args) finally: try: proc.wait(timeout=1) @@ -30,12 +30,12 @@ class RPCCase(unittest.TestCase): proc.kill() raise - def _blocking_echo(self): + def _blocking_echo(self, target): for attempt in range(100): time.sleep(.2) try: remote = pc_rpc.Client(test_address, test_port, - "test") + target) except ConnectionRefusedError: pass else: @@ -50,14 +50,17 @@ class RPCCase(unittest.TestCase): remote.close_rpc() def test_blocking_echo(self): - self._run_server_and_test(self._blocking_echo) + self._run_server_and_test(self._blocking_echo, "test") - async def _asyncio_echo(self): + def test_blocking_echo_autotarget(self): + self._run_server_and_test(self._blocking_echo, pc_rpc.AutoTarget) + + async def _asyncio_echo(self, target): remote = pc_rpc.AsyncioClient() for attempt in range(100): await asyncio.sleep(.2) try: - await remote.connect_rpc(test_address, test_port, "test") + await remote.connect_rpc(test_address, test_port, target) except ConnectionRefusedError: pass else: @@ -71,16 +74,19 @@ class RPCCase(unittest.TestCase): finally: remote.close_rpc() - def _loop_asyncio_echo(self): + def _loop_asyncio_echo(self, target): loop = asyncio.new_event_loop() asyncio.set_event_loop(loop) try: - loop.run_until_complete(self._asyncio_echo()) + loop.run_until_complete(self._asyncio_echo(target)) finally: loop.close() def test_asyncio_echo(self): - self._run_server_and_test(self._loop_asyncio_echo) + self._run_server_and_test(self._loop_asyncio_echo, "test") + + def test_asyncio_echo_autotarget(self): + self._run_server_and_test(self._loop_asyncio_echo, pc_rpc.AutoTarget) class FireAndForgetCase(unittest.TestCase):