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"]
