async def pub_handler(request): try: cmd = await request.json() except json.decoder.JSONDecodeError: raise aiohttp.web.HTTPBadRequest(body=b"Body must be json.") # Check cmd is a dict. if type(cmd) != dict: raise aiohttp.web.HTTPBadRequest(body=b"Body must be a json object.") # Check for required fields. if "container" not in cmd: raise aiohttp.web.HTTPBadRequest( body=b"Missing required 'container' field.") if "topic" not in cmd: raise aiohttp.web.HTTPBadRequest( body=b"Missing required 'topic' field.") # Fill optional fields. cmd["packet"] = cmd.get("packet", {}) headers = cmd["packet"].get("headers", []) payload = cmd["packet"].get("payload", "") # Find the absolute topic. tm = a0.TopicManager(container=cmd["container"]) topic = tm.publisher_topic(cmd["topic"]) # Perform requested action. p = a0.Publisher(topic) p.pub(headers, base64.b64decode(payload)) return aiohttp.web.Response(text="success")
async def sub_wshandler(request): ws = aiohttp.web.WebSocketResponse() await ws.prepare(request) msg = await ws.receive() cmd = json.loads(msg.data) tm = a0.TopicManager(container="api", subscriber_aliases={"topic": cmd}) init_ = { "OLDEST": a0.INIT_OLDEST, "MOST_RECENT": a0.INIT_MOST_RECENT, "AWAIT_NEW": a0.INIT_AWAIT_NEW, }[cmd["init"]] iter_ = {"NEXT": a0.ITER_NEXT, "NEWEST": a0.ITER_NEWEST}[cmd["iter"]] scheduler = cmd.get("scheduler", "IMMEDIATE") async for pkt in a0.aio_sub(tm.subscriber_topic("topic"), init_, iter_): await ws.send_json({ "headers": pkt.headers, "payload": base64.b64encode(pkt.payload).decode("utf-8"), }) if scheduler == "IMMEDIATE": pass elif scheduler == "ON_ACK": await ws.receive()
async def test_pub(sandbox): await sandbox.WaitUntilStartedAsync(timeout=1.0) async with aiohttp.ClientSession() as session: endpoint = "http://localhost:24880/api/pub" pub_data = { "container": "aaa", "topic": "bbb", "packet": { "payload": base64.b64encode(b"Hello, World!").decode("utf-8"), }, } # Normal publish. async with session.post(endpoint, data=json.dumps(pub_data)) as resp: assert resp.status == 200 assert await resp.text() == "success" # Normal publish. pub_data["packet"]["payload"] = base64.b64encode( b"Goodbye, World!").decode("utf-8") async with session.post(endpoint, data=json.dumps(pub_data)) as resp: assert resp.status == 200 assert await resp.text() == "success" # Missing "container". pub_data.pop("container") async with session.post(endpoint, data=json.dumps(pub_data)) as resp: assert resp.status == 400 assert await resp.text() == "Missing required 'container' field." pub_data["container"] = "aaa" # Missing "topic". pub_data.pop("topic") async with session.post(endpoint, data=json.dumps(pub_data)) as resp: assert resp.status == 400 assert await resp.text() == "Missing required 'topic' field." pub_data["topic"] = "bbb" # Not JSON. async with session.post(endpoint, data="not json") as resp: assert resp.status == 400 assert await resp.text() == "Body must be json." # Not JSON object. async with session.post(endpoint, data=json.dumps("not object")) as resp: assert resp.status == 400 assert await resp.text() == "Body must be a json object." tm = a0.TopicManager({"container": "aaa"}) sub = a0.SubscriberSync(tm.publisher_topic("bbb"), a0.INIT_OLDEST, a0.ITER_NEXT) msgs = [] while sub.has_next(): msgs.append(sub.next().payload) assert len(msgs) == 2 assert msgs == [b"Hello, World!", b"Goodbye, World!"]
async def prpc_wshandler(request): ws = aiohttp.web.WebSocketResponse() await ws.prepare(request) msg = await ws.receive() try: cmd = json.loads(msg.data) except json.JSONDecodeError: await ws.close(message=b"Message must be json.") return # Check cmd is a dict. if type(cmd) != dict: await ws.close(message=b"Message must be a json object.") return # Fill optional fields. cmd["packet"] = cmd.get("packet", {}) headers = cmd["packet"].get("headers", []) payload = cmd["packet"].get("payload", "") scheduler = cmd.get("scheduler", "IMMEDIATE") tm = a0.TopicManager(container="api", prpc_client_aliases={"topic": cmd}) ns = types.SimpleNamespace() ns.loop = asyncio.get_event_loop() ns.q = asyncio.Queue() def prpc_callback(pkt_view, done): pkt = a0.Packet(pkt_view) ns.loop.call_soon_threadsafe(ns.q.put_nowait, (pkt, done)) prpc_client = a0.PrpcClient(tm.prpc_client_topic("topic")) req = a0.Packet(headers, base64.b64decode(payload)) prpc_client.connect(req, prpc_callback) while True: pkt, done = await ns.q.get() await ws.send_json({ "headers": pkt.headers, "payload": base64.b64encode(pkt.payload).decode("utf-8"), }) if done: break if scheduler == "IMMEDIATE": pass elif scheduler == "ON_ACK": await ws.receive()
async def pub_wshandler(request): ws = aiohttp.web.WebSocketResponse() await ws.prepare(request) handshake_completed = False async for msg in ws: if msg.type != aiohttp.WSMsgType.TEXT: break try: cmd = json.loads(msg.data) except json.JSONDecodeError: await ws.close(message=b"Message must be json.") return # Check cmd is a dict. if type(cmd) != dict: await ws.close(message=b"Message must be a json object.") return if not handshake_completed: # Check for required fields. if "container" not in cmd: await ws.close(message=b"Missing required 'container' field.") return if "topic" not in cmd: await ws.close(message=b"Missing required 'topic' field.") return # Create a publisher on the absolute topic. tm = a0.TopicManager(container=cmd["container"]) publisher = a0.Publisher(tm.publisher_topic(cmd["topic"])) handshake_completed = True continue # Fill optional fields. cmd["packet"] = cmd.get("packet", {}) headers = cmd["packet"].get("headers", []) payload = cmd["packet"].get("payload", "") publisher.pub(headers, base64.b64decode(payload))
async def rpc_handler(request): try: cmd = await request.json() except json.decoder.JSONDecodeError: raise aiohttp.web.HTTPBadRequest(body=b"Body must be json.") # Check cmd is a dict. if type(cmd) != dict: raise aiohttp.web.HTTPBadRequest(body=b"Body must be a json object.") # Check for required fields. if "container" not in cmd: raise aiohttp.web.HTTPBadRequest( body=b"Missing required 'container' field.") if "topic" not in cmd: raise aiohttp.web.HTTPBadRequest( body=b"Missing required 'topic' field.") # Fill optional fields. cmd["packet"] = cmd.get("packet", {}) headers = cmd["packet"].get("headers", []) payload = cmd["packet"].get("payload", "") # Find the absolute topic. tm = a0.TopicManager(container="api", rpc_client_aliases={ "topic": cmd, }) topic = tm.rpc_client_topic("topic") # Perform requested action. client = a0.AioRpcClient(topic) resp = await client.send(a0.Packet(headers, base64.b64decode(payload))) return aiohttp.web.json_response({ "headers": resp.headers, "payload": base64.b64encode(resp.payload).decode("utf-8"), })
async def test_ws_pub(sandbox): await sandbox.WaitUntilStartedAsync(timeout=1.0) async with aiohttp.ClientSession() as session: endpoint = "ws://localhost:24880/wsapi/pub" #################### # Basic test case. # #################### handshake_data = { "container": "aaa", "topic": "bbb", } ws = await session.ws_connect(endpoint) await ws.send_str(json.dumps(handshake_data)) async def send_payload(payload): payload_utf8 = payload.encode("utf-8") payload_b64 = base64.b64encode(payload_utf8).decode("utf-8") await ws.send_str(json.dumps({"packet": {"payload": payload_b64}})) await asyncio.gather( send_payload("message_0"), send_payload("message_1"), send_payload("message_2"), ) await asyncio.sleep(0.1) tm = a0.TopicManager({"container": "aaa"}) sub = a0.SubscriberSync(tm.publisher_topic("bbb"), a0.INIT_OLDEST, a0.ITER_NEXT) msgs = set() while sub.has_next(): msgs.add(sub.next().payload) assert len(msgs) == 3 assert msgs == set([b"message_0", b"message_1", b"message_2"]) await ws.close() ################## # Bad handshake. # ################## # Not JSON. ws = await session.ws_connect(endpoint) await ws.send_str("not json") reply = await ws.receive() assert reply.type in [ aiohttp.WSMsgType.CLOSE, aiohttp.WSMsgType.CLOSED ] assert reply.extra == "Message must be json." assert ws.closed # Not JSON object. ws = await session.ws_connect(endpoint) await ws.send_str(json.dumps("not object")) reply = await ws.receive() assert reply.type in [ aiohttp.WSMsgType.CLOSE, aiohttp.WSMsgType.CLOSED ] assert reply.extra == "Message must be a json object." assert ws.closed # Missing "container". ws = await session.ws_connect(endpoint) await ws.send_str(json.dumps({"foo": "bar"})) reply = await ws.receive() assert reply.type in [ aiohttp.WSMsgType.CLOSE, aiohttp.WSMsgType.CLOSED ] assert reply.extra == "Missing required 'container' field." assert ws.closed # Missing "topic". ws = await session.ws_connect(endpoint) await ws.send_str(json.dumps({"container": "aaa"})) reply = await ws.receive() assert reply.type in [ aiohttp.WSMsgType.CLOSE, aiohttp.WSMsgType.CLOSED ] assert reply.extra == "Missing required 'topic' field." assert ws.closed ############### # Bad packet. # ############### # Not JSON. ws = await session.ws_connect(endpoint) await ws.send_str(json.dumps(handshake_data)) await ws.send_str("not json") reply = await ws.receive() assert reply.type in [ aiohttp.WSMsgType.CLOSE, aiohttp.WSMsgType.CLOSED ] assert reply.extra == "Message must be json." assert ws.closed # Not JSON object. ws = await session.ws_connect(endpoint) await ws.send_str(json.dumps(handshake_data)) await ws.send_str(json.dumps("not object")) reply = await ws.receive() assert reply.type in [ aiohttp.WSMsgType.CLOSE, aiohttp.WSMsgType.CLOSED ] assert reply.extra == "Message must be a json object." assert ws.closed
async def test_rpc(sandbox): await sandbox.WaitUntilStartedAsync(timeout=1.0) ns = types.SimpleNamespace() ns.collected_requests = [] def on_request(req): ns.collected_requests.append(req.pkt.payload.decode("utf-8")) req.reply(f"success_{len(ns.collected_requests)}") tm = a0.TopicManager({"container": "aaa"}) topic = tm.rpc_server_topic("bbb") server = a0.RpcServer(topic, on_request, None) # noqa: F841 async with aiohttp.ClientSession() as session: endpoint = "http://localhost:24880/api/rpc" rpc_data = { "container": "aaa", "topic": "bbb", "packet": { "payload": "", }, } # Normal request. rpc_data["packet"]["payload"] = base64.b64encode(b"request_0").decode( "utf-8") async with session.post(endpoint, data=json.dumps(rpc_data)) as resp: assert resp.status == 200 resp_pkt = await resp.json() assert base64.b64decode(resp_pkt.get("payload", "")) == b"success_1" # Normal request. rpc_data["packet"]["payload"] = base64.b64encode(b"request_1").decode( "utf-8") async with session.post(endpoint, data=json.dumps(rpc_data)) as resp: assert resp.status == 200 resp_pkt = await resp.json() assert base64.b64decode(resp_pkt.get("payload", "")) == b"success_2" # Missing "container". rpc_data.pop("container") async with session.post(endpoint, data=json.dumps(rpc_data)) as resp: assert resp.status == 400 assert await resp.text() == "Missing required 'container' field." rpc_data["container"] = "aaa" # Missing "topic". rpc_data.pop("topic") async with session.post(endpoint, data=json.dumps(rpc_data)) as resp: assert resp.status == 400 assert await resp.text() == "Missing required 'topic' field." rpc_data["topic"] = "bbb" # Not JSON. async with session.post(endpoint, data="not json") as resp: assert resp.status == 400 assert await resp.text() == "Body must be json." # Not JSON object. async with session.post(endpoint, data=json.dumps("not object")) as resp: assert resp.status == 400 assert await resp.text() == "Body must be a json object." assert ns.collected_requests == ["request_0", "request_1"]