async def handle_single_request(self, raw: Optional[Message]):
if raw is None:
return
msg, buffers = raw
if "id" not in msg:
await self.write(Message({
'id': None,
'type': 'error',
'error': 'ID missing from request-response.'
}))
return
id = msg['id']
try:
if "type" not in msg:
raise ReqRespServerException("Type missing from request.")
type = msg["type"]
if type == "request":
result: Message = await self.handle_request(id, msg, buffers)
await self.write(result)
else:
raise ReqRespServerException('Type missing from frame.')
except ReqRespServerException as e:
await self.write(Message({'id': id, 'type': 'error', 'error': str(e)}))
async def on_render(
self,
frame: int,
format: Optional[list] = None,
planes: Optional[Union[List[int], int]] = None) -> Message:
frame_inst = self.clip[frame]
async with frame_inst:
if planes is None:
return Message({'size': list(frame_inst.size)}, [])
format: RawFormat = RawFormat.from_json(format)
if not (await frame_inst.can_render(format)):
return Message({'size': None})
if not isinstance(planes, (list, tuple)):
planes: List[int] = [planes]
buffers = [
bytearray(format.get_plane_size(p, frame_inst.size))
for p in planes
]
used_buffers = await gather(
*(frame_inst.render_into(buf, p, format, 0)
for p, buf in enumerate(buffers)))
return Message(
{'size': list(frame_inst.size)},
[buf[:sz] for buf, sz in zip(buffers, used_buffers)])
async def test_read_stream(self):
rsock, wsock = socket.socketpair()
rt = None
try:
loop = get_running_loop()
rt, proto = await loop.create_connection(YuunoProtocol, sock=rsock)
stream = ConnectionInputStream(proto)
wsock.sendall(
ByteOutputStream.write_message(Message({"test": id(self)},
[])))
async with stream:
msg = await wait_for(stream.read(), 5)
self.assertEqual(msg, Message({"test": id(self)}, []))
wsock.shutdown(socket.SHUT_RDWR)
await wait_for(stream.read(), 5)
self.assertFalse(stream.acquired)
finally:
if rt is not None:
rt.close()
wsock.close()
def test_read_text_only(self):
parser = bytes_protocol()
self.assertEqual(
list(parser.feed(b"\x00\x00\x00\x01\x00\x00\x00\x02{}")),
[Message({}, [])])
self.assertEqual(
list(parser.feed(b'\x00\x00\x00\x01\x00\x00\x00\x0a{"test":1}')),
[Message({"test": 1}, [])])
def test_write_message_buffer(self):
self.assertEqual(
ByteOutputStream.write_message(Message({}, [b"ab"])),
b'\x00\x00\x00\x02\x00\x00\x00\x02\x00\x00\x00\x02{}ab')
self.assertEqual(
ByteOutputStream.write_message(Message({}, [b"ab", b"cd"])),
b'\x00\x00\x00\x03\x00\x00\x00\x02\x00\x00\x00\x02\x00\x00\x00\x02{}abcd'
)
async def on_get_config(self, key: str):
result = await self.script.get_config(key, NOT_SET)
if result is NOT_SET:
return Message({"type": "unset", "value": None}, [])
vtype, converter = TYPE_MAP_SEND[type(result)]
value, buffers = converter(result)
return Message({"type": vtype, "value": value}, buffers)
def test_read_with_buffer(self):
parser = bytes_protocol()
self.assertEqual(
list(
parser.feed(
b'\x00\x00\x00\x02\x00\x00\x00\x02\x00\x00\x00\x02{}ab')),
[Message({}, [b"ab"])])
self.assertEqual(
list(
parser.feed(
b'\x00\x00\x00\x03\x00\x00\x00\x02\x00\x00\x00\x02\x00\x00\x00\x02{}abcd'
)), [Message({}, [b"ab", b"cd"])])
def test_read_both(self):
parser = bytes_protocol()
self.assertEqual(
list(
parser.feed(
b'\x00\x00\x00\x03\x00\x00\x00\x0a\x00\x00\x00\x02\x00\x00\x00\x02{"test":1}abcd'
)), [Message({"test": 1}, [b"ab", b"cd"])])
async def test_write_stream(self):
rsock, wsock = socket.socketpair()
wt = FakeClosable()
rs_w = FakeClosable()
try:
rs_r, rs_w = await open_connection(sock=rsock)
wt, wp = await get_running_loop().create_connection(YuunoProtocol,
sock=wsock)
wpos = ConnectionOutputStream(wp)
async with wpos:
msg = Message({"test": id(self)}, [])
msg_raw = ByteOutputStream.write_message(msg)
plength = len(msg_raw)
await wait_for(wpos.write(msg), 5)
data = await wait_for(rs_r.readexactly(plength), 5)
self.assertEqual(data, msg_raw)
await wait_for(rs_r.read(1), 5)
self.assertTrue(rs_r.at_eof())
finally:
rs_w.close()
wt.close()
async def handle_request(self, id, msg: JSON, buffers: List[bytes]):
try:
if "method" not in msg:
raise ReqRespServerException("No method given.")
method = msg['method']
func = getattr(self, "on_" + method.lower(), None)
if func is None:
raise ReqRespServerException("Unknown method.")
additional = {}
if len(buffers) > 0:
additional["_buffers"] = buffers
result: Union[Awaitable[Message], Message] = func(**additional, **msg.get('params', {}))
if inspect.isawaitable(result):
result = await result
result: Message = cast('Message', result)
return Message({"id": id, "type": "response", "result": result.values}, result.blobs)
except ReqRespServerException:
raise
except Exception as e:
exc = ''.join(format_exception(type(e), e, e.__traceback__))
raise ReqRespServerException("An error occured while executing the function:\n" + exc) from None
async def _delivered(self, raw: Optional[Message]):
if raw is None:
waiters = self._waiters
self._waiters = None
for v in waiters.values():
v.cancel()
return
msg, buffers = raw
if "id" not in msg:
return
id = msg["id"]
future: Future[Message] = self._waiters.pop(id, None)
if future is None:
return
if "type" not in msg:
future.set_exception(ReqRespServerException("Missing type in server response."))
return
type = msg["type"]
if type == "error":
future.set_exception(CallFailed(msg.get("error", "The requested funtion failed.")))
else:
future.set_result(Message(msg["result"], buffers))
async def write(self, message: Message) -> NoReturn:
await self.channel.multiplexer.write(
Message(
{
'target': self.channel.name,
'type': 'message',
'payload': message.values
}, message.blobs))
async def test_echo(self):
c1, c2 = pipe_bidi()
srv = MockReqRespServer(c1)
cli = MockReqRespClient(c2)
async with srv, cli:
result: Message = await wait_for(cli.echo(), 5)
self.assertEqual(Message({}, []), result)
async def test_connection_receive_twice(self):
pc = Connection(*pipe())
m = Message({}, [])
async with pc:
await pc.write(m)
self.assertIs(m, await pc.read())
with self.assertRaises(TimeoutError):
await wait_for(pc.read(), 1)
async def test_pipe_read(self):
r, w = pipe()
with r, w:
fis = FileInputStream(r)
async with timeout_context(fis, 5):
msg = Message({"test": id(self)}, [])
w.write(ByteOutputStream.write_message(msg))
w.flush()
self.assertEqual(await wait_for(fis.read(), 5), msg)
async def on_run(self,
encoding: Optional[str] = None,
_buffers: Sequence[bytes] = (b"", )):
code = _buffers[0]
if encoding is not None:
code = code.decode(encoding)
await self.script.run(code)
return Message({}, [])
async def on_set_config(self,
key: str,
value: ConfigTypes,
type: str,
_buffers: List[bytes] = ()):
converter: ConverterTypeRecv = TYPE_MAP_RECV[type]
value = converter(value, _buffers)
await self.script.set_config(key, value)
return Message({}, [])
async def _call(self, funcname: str, buffers: List[bytes], params: Mapping[str, JSON]) -> Awaitable[Message]:
cid = self._current_id
self._current_id += 1
future = get_running_loop().create_future()
self._waiters[cid] = future
await self.write(Message({"method": funcname, "type": "request", "id": cid, "params": params}, buffers))
return (await future)
async def close(self) -> NoReturn:
self.channel.multiplexer.streams.pop(self.channel.name, None)
if self.channel.ingress is not None and self.channel.ingress.output is not None:
await self.channel.ingress.close()
if self.channel.multiplexer.acquired:
await self.channel.multiplexer.write(
Message({
'target': self.channel.name,
'type': 'close'
}))
async def test_receiving_messages(self):
m1 = Message({}, [])
m2 = Message({}, [])
messages = []
finished = Event()
async def _r(m: Message):
messages.append(m)
if len(m) == 2:
finished.set()
pipe_r, pipe_w = pipe()
reader = ReaderTask(pipe_r, _r)
async with reader, pipe_w:
await pipe_w.write(m1)
await pipe_w.write(m2)
await wait_for(finished.wait(), 5)
self.assertIs(m1, messages[0])
self.assertIs(m2, messages[1])
async def _delivered(self, raw: Optional[Message]) -> None:
if not self.acquired:
return
if raw is None:
self._shutdown.set()
return
msg, buffers = raw
connection = msg.get("target", "")
type = msg.get("type", "message")
if type in ("close", "illegal"):
reader = self.streams.pop(connection, None)
if reader is None:
return
await reader.deliver(None)
return
if connection not in self.streams:
await self.write(
Message({
"target": connection,
"type": "close",
"payload": {}
}))
return
if "payload" not in msg:
await self.write(
Message({
"target": connection,
"type": "illegal",
"payload": {}
}))
return
await self.streams[connection].deliver(Message(msg["payload"],
buffers))
return
async def test_invalid_packet__unset_method(self):
c1, c2 = pipe_bidi()
srv = MockReqRespServer(c1)
async with srv, c2:
await c2.write(Message({'id': 0, 'type': 'request', 'params': {}}))
msg: Message = await wait_for(c2.read(), 5)
self.assertIsInstance(msg.values, dict)
self.assertIn('type', msg.values)
self.assertEqual('error', msg.values['type'])
self.assertIn('id', msg.values)
self.assertEqual(0, msg.values['id'])
self.assertIn('error', msg.values)
async def on_create_script(self, channel_name: str, params: Dict[str,
Any]):
conn = self.multiplexer.connect(channel_name)
register(self, conn)
await conn.acquire()
script = await self.provider.get(**params)
server = RemoteScriptServer(script, conn)
register(conn, server)
await server.acquire()
self.connections[channel_name] = conn
return Message({}, [])
async def test_multiplexer_reject(self):
c_faked_networking, c_multiplexed = pipe_bidi()
muliplexer = Multiplexer(c_multiplexed)
async with c_faked_networking, muliplexer:
await c_faked_networking.write(
Message({
'target': 'non-existent',
'payload': {}
}))
msg: Message = await wait_for(c_faked_networking.read(), 5)
self.assertIsInstance(msg.values, dict)
self.assertIn('type', msg.values)
self.assertEqual('close', msg.values['type'])
async def on_register_clip(self, channel: str, name: str):
conn = self.multiplexer.connect(channel)
register(self, conn)
await conn.acquire()
clips = await self.script.retrieve_clips()
clip = clips[name]
server = ClipServer(clip, conn)
register(conn, server)
await server.acquire()
self.connections[channel] = conn
return Message({}, [])
async def test_connection_read_with_wait(self):
pc = Connection(*pipe())
m = Message({}, [])
async def _concurrent():
t0 = time.time()
self.assertIs(m, await pc.read())
t1 = time.time()
self.assertGreaterEqual(t1-t0, 1)
async with pc:
task = get_running_loop().create_task(_concurrent())
await sleep(1)
await pc.write(m)
await task
async def test_multiplexer_illegal(self):
c_faked_networking, c_multiplexed = pipe_bidi()
muliplexer = Multiplexer(c_multiplexed)
async with c_faked_networking, muliplexer:
async with muliplexer.connect('existing') as f:
await c_faked_networking.write(
Message({
'target': 'existing',
'type': 'message'
}))
msg: Message = await wait_for(c_faked_networking.read(), 5)
self.assertIsInstance(msg.values, dict)
self.assertIn('type', msg.values)
self.assertEqual('illegal', msg.values['type'])
async def test_multiplexer_delivered(self):
value = {'v': 'test'}
c_faked_networking, c_multiplexed = pipe_bidi()
muliplexer = Multiplexer(c_multiplexed)
async with c_faked_networking, muliplexer:
async with muliplexer.connect('existing') as f:
await c_faked_networking.write(
Message({
'target': 'existing',
'type': 'message',
'payload': value
}))
msg: Message = await wait_for(f.read(), 5)
self.assertIs(value, msg.values)
async def test_connection_close_with_wait(self):
pc = Connection(*pipe())
m = Message({}, [])
async def _concurrent():
t0 = time.time()
with self.assertRaises(TimeoutError):
await pc.read()
t1 = time.time()
self.assertGreaterEqual(t1-t0, 1)
async with pc:
task = get_running_loop().create_task(_concurrent())
await sleep(1.2)
await pc.close()
await task
async def test_multiplexer_send(self):
values = {}
c_faked_networking, c_multiplexed = pipe_bidi()
muliplexer = Multiplexer(c_multiplexed)
async with c_faked_networking, muliplexer:
async with muliplexer.connect('existing') as f:
await f.write(Message(values, []))
msg: Message = await wait_for(c_faked_networking.read(), 5)
self.assertIsInstance(msg.values, dict)
self.assertIn('target', msg.values)
self.assertEqual('existing', msg.values['target'])
self.assertIn('type', msg.values)
self.assertEqual('message', msg.values['type'])
self.assertIn('payload', msg.values)
self.assertIs(values, msg.values['payload'])