139 lines
4.3 KiB
Python
139 lines
4.3 KiB
Python
from __future__ import annotations
|
|
|
|
import argparse
|
|
import asyncio
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
|
|
|
|
from toki_socket.communicator import type_name_of
|
|
from toki_socket.packets.message_common_pb2 import TestData
|
|
from toki_socket.tcp_client import connect_tcp
|
|
from toki_socket.ws_client import connect_ws
|
|
|
|
HOST = "127.0.0.1"
|
|
WS_PATH = "/"
|
|
CONNECT_WINDOW = 3.0
|
|
REQUEST_WINDOW = 2.0
|
|
|
|
|
|
def parser_map():
|
|
return {type_name_of(TestData): TestData.FromString}
|
|
|
|
|
|
async def main() -> None:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--mode", choices=["tcp", "ws"], default="tcp")
|
|
parser.add_argument("--port", type=int, required=True)
|
|
parser.add_argument("--phase", choices=["send-push", "requests"], default="send-push")
|
|
args = parser.parse_args()
|
|
|
|
print(f"INFO typeName python={type_name_of(TestData)}")
|
|
|
|
try:
|
|
client = await dial_with_retry(args.mode, args.port)
|
|
except Exception as exc:
|
|
fail("setup", str(exc))
|
|
raise SystemExit(1) from exc
|
|
|
|
try:
|
|
if args.phase == "send-push":
|
|
ok = await run_send_push(client)
|
|
else:
|
|
ok = await run_requests(client)
|
|
finally:
|
|
await client.close()
|
|
|
|
if not ok:
|
|
raise SystemExit(1)
|
|
|
|
|
|
async def dial_with_retry(mode: str, port: int):
|
|
deadline = asyncio.get_running_loop().time() + CONNECT_WINDOW
|
|
last_error: Exception | None = None
|
|
while asyncio.get_running_loop().time() < deadline:
|
|
try:
|
|
return await asyncio.wait_for(dial(mode, port), 0.3)
|
|
except Exception as exc:
|
|
last_error = exc
|
|
await asyncio.sleep(0.1)
|
|
raise TimeoutError(f"connect {mode}:{port} timed out: {last_error}")
|
|
|
|
|
|
async def dial(mode: str, port: int):
|
|
if mode == "tcp":
|
|
return await connect_tcp(HOST, port, 0, 0, parser_map())
|
|
if mode == "ws":
|
|
return await connect_ws(HOST, port, WS_PATH, 0, 0, parser_map())
|
|
raise ValueError(f"unknown mode {mode!r}")
|
|
|
|
|
|
async def run_send_push(client) -> bool:
|
|
loop = asyncio.get_running_loop()
|
|
push = loop.create_future()
|
|
client.communicator.add_listener(
|
|
type_name_of(TestData),
|
|
lambda message: push.set_result(message) if not push.done() else None,
|
|
)
|
|
|
|
try:
|
|
await client.communicator.send(TestData(index=101, message="fire from python client"))
|
|
passed("1", "fire-and-forget sent")
|
|
|
|
message = await asyncio.wait_for(push, REQUEST_WINDOW)
|
|
if message.index != 200 or message.message != "push from kotlin server":
|
|
fail("2", f"unexpected push index={message.index} message={message.message!r}")
|
|
return False
|
|
passed("2", "received push from kotlin server")
|
|
return True
|
|
except Exception as exc:
|
|
fail("2", str(exc))
|
|
return False
|
|
|
|
|
|
async def run_requests(client) -> bool:
|
|
try:
|
|
single = await client.communicator.send_request(
|
|
TestData(index=21, message="single request from python"),
|
|
TestData,
|
|
timeout=REQUEST_WINDOW,
|
|
)
|
|
if single.index != 42 or single.message != "echo: single request from python":
|
|
fail("3", f"unexpected response index={single.index} message={single.message!r}")
|
|
return False
|
|
passed("3", "single request response matched")
|
|
except Exception as exc:
|
|
fail("3", str(exc))
|
|
return False
|
|
|
|
async def request_one(i: int) -> None:
|
|
index = 30 + i
|
|
message = f"multi request {i} from python"
|
|
res = await client.communicator.send_request(
|
|
TestData(index=index, message=message),
|
|
TestData,
|
|
timeout=REQUEST_WINDOW,
|
|
)
|
|
if res.index != index * 2 or res.message != "echo: " + message:
|
|
raise AssertionError(f"request {i} got index={res.index} message={res.message!r}")
|
|
|
|
try:
|
|
await asyncio.gather(*(request_one(i) for i in range(5)))
|
|
passed("4", "concurrent request responses matched")
|
|
return True
|
|
except Exception as exc:
|
|
fail("4", str(exc))
|
|
return False
|
|
|
|
|
|
def passed(scenario: str, detail: str) -> None:
|
|
print(f"PASS scenario={scenario} detail={detail}")
|
|
|
|
|
|
def fail(scenario: str, error: str) -> None:
|
|
print(f"FAIL scenario={scenario} error={error}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
asyncio.run(main())
|