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'])