artiq/artiq/test/test_pc_rpc.py

133 lines
3.9 KiB
Python
Raw Normal View History

2014-10-25 17:25:30 +08:00
import unittest
import sys
import subprocess
import asyncio
import time
import numpy as np
2015-06-15 12:03:34 +08:00
from artiq.protocols import pc_rpc, fire_and_forget
2014-10-25 17:25:30 +08:00
test_address = "::1"
test_port = 7777
2015-01-05 16:07:12 +08:00
test_object = [5, 2.1, None, True, False,
{"a": 5, 2: np.linspace(0, 10, 1)},
(4, 5), (10,), "ab\nx\"'"]
2014-10-25 17:25:30 +08:00
class RPCCase(unittest.TestCase):
2015-10-19 20:20:53 +08:00
def _run_server_and_test(self, test, *args):
2014-10-25 17:25:30 +08:00
# running this file outside of unittest starts the echo server
with subprocess.Popen([sys.executable,
sys.modules[__name__].__file__]) as proc:
try:
2015-10-19 20:20:53 +08:00
test(*args)
2014-10-25 17:25:30 +08:00
finally:
try:
proc.wait(timeout=1)
except subprocess.TimeoutExpired:
proc.kill()
raise
2015-01-05 16:07:12 +08:00
2015-10-19 20:20:53 +08:00
def _blocking_echo(self, target):
2015-01-05 16:07:12 +08:00
for attempt in range(100):
time.sleep(.2)
try:
remote = pc_rpc.Client(test_address, test_port,
2015-10-19 20:20:53 +08:00
target)
2015-01-05 16:07:12 +08:00
except ConnectionRefusedError:
pass
else:
break
try:
test_object_back = remote.echo(test_object)
self.assertEqual(test_object, test_object_back)
test_object_back = remote.async_echo(test_object)
self.assertEqual(test_object, test_object_back)
with self.assertRaises(AttributeError):
2015-01-05 16:07:12 +08:00
remote.non_existing_method()
2015-08-17 23:17:13 +08:00
remote.terminate()
2015-01-05 16:07:12 +08:00
finally:
remote.close_rpc()
def test_blocking_echo(self):
2015-10-19 20:20:53 +08:00
self._run_server_and_test(self._blocking_echo, "test")
2015-01-05 16:07:12 +08:00
2015-10-19 20:20:53 +08:00
def test_blocking_echo_autotarget(self):
self._run_server_and_test(self._blocking_echo, pc_rpc.AutoTarget)
async def _asyncio_echo(self, target):
2015-01-05 16:07:12 +08:00
remote = pc_rpc.AsyncioClient()
for attempt in range(100):
2015-10-03 19:28:57 +08:00
await asyncio.sleep(.2)
2015-01-05 16:07:12 +08:00
try:
2015-10-19 20:20:53 +08:00
await remote.connect_rpc(test_address, test_port, target)
2015-01-05 16:07:12 +08:00
except ConnectionRefusedError:
pass
else:
break
try:
2015-10-03 19:28:57 +08:00
test_object_back = await remote.echo(test_object)
2014-10-25 17:25:30 +08:00
self.assertEqual(test_object, test_object_back)
test_object_back = await remote.async_echo(test_object)
self.assertEqual(test_object, test_object_back)
with self.assertRaises(AttributeError):
2015-10-03 19:28:57 +08:00
await remote.non_existing_method()
await remote.terminate()
2015-01-05 16:07:12 +08:00
finally:
remote.close_rpc()
2015-10-19 20:20:53 +08:00
def _loop_asyncio_echo(self, target):
loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)
try:
2015-10-19 20:20:53 +08:00
loop.run_until_complete(self._asyncio_echo(target))
finally:
loop.close()
2015-01-05 16:07:12 +08:00
def test_asyncio_echo(self):
2015-10-19 20:20:53 +08:00
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)
2014-10-25 17:25:30 +08:00
2015-06-15 12:03:34 +08:00
class FireAndForgetCase(unittest.TestCase):
def _set_ok(self):
self.ok = True
def test_fire_and_forget(self):
self.ok = False
p = fire_and_forget.FFProxy(self)
p._set_ok()
p.ff_join()
self.assertTrue(self.ok)
class Echo:
2014-10-25 17:25:30 +08:00
def echo(self, x):
return x
async def async_echo(self, x):
await asyncio.sleep(0.01)
return x
2014-10-25 17:25:30 +08:00
def run_server():
loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)
2014-10-25 17:25:30 +08:00
try:
echo = Echo()
2015-08-17 23:17:13 +08:00
server = pc_rpc.Server({"test": echo}, builtin_terminate=True)
2014-10-25 17:25:30 +08:00
loop.run_until_complete(server.start(test_address, test_port))
try:
2015-08-17 23:17:13 +08:00
loop.run_until_complete(server.wait_terminate())
2014-10-25 17:25:30 +08:00
finally:
loop.run_until_complete(server.stop())
finally:
loop.close()
if __name__ == "__main__":
run_server()