generated from mononoki/pyproj
[+!] all but small
some small modifications, tests added
This commit is contained in:
176
examples/client.py
Normal file
176
examples/client.py
Normal file
@@ -0,0 +1,176 @@
|
||||
"""Small Artex v1 client, shared by the manual run and server tests."""
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
from contextlib import suppress
|
||||
import json
|
||||
|
||||
|
||||
class ProtocolError(RuntimeError):
|
||||
pass
|
||||
|
||||
|
||||
class Client:
|
||||
def __init__(
|
||||
self, name="test-rover", host="127.0.0.1", port=9000,
|
||||
timeout=5.0, node_type="test", inputs=(), outputs=("telemetry",),
|
||||
trace=None,
|
||||
):
|
||||
self.name = name
|
||||
self.host = host
|
||||
self.port = port
|
||||
self.timeout = timeout
|
||||
self.node_type = node_type
|
||||
self.inputs = list(inputs)
|
||||
self.outputs = list(outputs)
|
||||
self.trace = trace
|
||||
self.reader = None
|
||||
self.writer = None
|
||||
self.uid = None
|
||||
self.upf = None
|
||||
self.pending = []
|
||||
|
||||
async def send(self, command, payload):
|
||||
if self.writer is None:
|
||||
raise ProtocolError("Client is not connected")
|
||||
packet = {"command": command, "payload": payload}
|
||||
self.writer.write((json.dumps(packet) + "\n").encode("utf-8"))
|
||||
await asyncio.wait_for(self.writer.drain(), self.timeout)
|
||||
if self.trace:
|
||||
self.trace("->", packet)
|
||||
|
||||
async def _read(self):
|
||||
raw = await self.reader.readline()
|
||||
if not raw:
|
||||
raise ProtocolError("Server closed the connection")
|
||||
try:
|
||||
packet = json.loads(raw)
|
||||
except (ValueError, UnicodeDecodeError) as error:
|
||||
raise ProtocolError("Server returned invalid JSON") from error
|
||||
if not isinstance(packet, dict) or not isinstance(packet.get("command"), str):
|
||||
raise ProtocolError("Server packet must contain a string command")
|
||||
if self.trace:
|
||||
self.trace("<-", packet)
|
||||
return packet
|
||||
|
||||
async def recv(self):
|
||||
if self.pending:
|
||||
return self.pending.pop(0)
|
||||
return await asyncio.wait_for(self._read(), self.timeout)
|
||||
|
||||
async def _expect(self, command, allow_data=False):
|
||||
# One deadline covers the whole wait, including interleaved data packets.
|
||||
async with asyncio.timeout(self.timeout):
|
||||
while True:
|
||||
packet = await self._read()
|
||||
received = packet["command"]
|
||||
if received == command:
|
||||
return packet
|
||||
if allow_data and received not in {
|
||||
"token_ask", "handshake_ask", "heartbeat_answer",
|
||||
}:
|
||||
self.pending.append(packet)
|
||||
continue
|
||||
raise ProtocolError(f"Expected {command}, received {received}")
|
||||
|
||||
async def connect(self, old_token=None, old_upf=None):
|
||||
if (old_token is None) != (old_upf is None):
|
||||
raise ValueError("Recovery requires both old_token and old_upf")
|
||||
if self.writer is not None:
|
||||
raise ProtocolError("Client is already connected")
|
||||
self.pending.clear()
|
||||
self.uid = None
|
||||
self.upf = None
|
||||
self.reader, self.writer = await asyncio.wait_for(
|
||||
asyncio.open_connection(self.host, self.port), self.timeout,
|
||||
)
|
||||
try:
|
||||
await self.send("connection", {
|
||||
"name": self.name,
|
||||
"type": self.node_type,
|
||||
"in": self.inputs,
|
||||
"out": self.outputs,
|
||||
})
|
||||
token = await self._expect("token_ask")
|
||||
self.uid = token.get("payload")
|
||||
self.upf = token.get("upf")
|
||||
if type(self.uid) is not int or type(self.upf) is not int:
|
||||
raise ProtocolError("token_ask must contain integer payload and upf")
|
||||
|
||||
await self.send("token_answer", {"token": old_token, "upf": old_upf})
|
||||
await self._expect("handshake_ask")
|
||||
await self.send("handshake_answer", {"status": True, "state": None})
|
||||
# The heartbeat reply confirms that the server reached the air state.
|
||||
await self.heartbeat()
|
||||
except BaseException:
|
||||
with suppress(OSError, TimeoutError):
|
||||
await self.close()
|
||||
raise
|
||||
return self
|
||||
|
||||
async def heartbeat(self):
|
||||
await self.send("heartbeat_ask", {})
|
||||
return await self._expect("heartbeat_answer", allow_data=True)
|
||||
|
||||
async def telemetry(self, temperature=42.0, battery=87):
|
||||
await self.send("telemetry", {"temperature": temperature, "battery": battery})
|
||||
|
||||
async def close(self):
|
||||
writer = self.writer
|
||||
self.writer = None
|
||||
self.reader = None
|
||||
if writer is not None:
|
||||
writer.close()
|
||||
with suppress(ConnectionError):
|
||||
await asyncio.wait_for(writer.wait_closed(), self.timeout)
|
||||
|
||||
|
||||
def print_packet(direction, packet):
|
||||
print(direction, json.dumps(packet, ensure_ascii=False), flush=True)
|
||||
|
||||
|
||||
async def run(args):
|
||||
client = Client(
|
||||
args.name, args.host, args.port, args.timeout,
|
||||
node_type=args.node_type, trace=print_packet,
|
||||
)
|
||||
try:
|
||||
await client.connect(args.old_token, args.old_upf)
|
||||
print(f"Connected: uid={client.uid}, upf={client.upf}", flush=True)
|
||||
if args.old_token is not None:
|
||||
print("Recovery requested; the current protocol has no recovery-result reply.")
|
||||
for index in range(args.count):
|
||||
await client.telemetry()
|
||||
await client.heartbeat()
|
||||
if index + 1 < args.count:
|
||||
await asyncio.sleep(args.interval)
|
||||
finally:
|
||||
await client.close()
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("name", nargs="?", default="test-rover")
|
||||
parser.add_argument("--host", default="127.0.0.1")
|
||||
parser.add_argument("--port", type=int, default=9000)
|
||||
parser.add_argument("--type", dest="node_type", default="test")
|
||||
parser.add_argument("--timeout", type=float, default=5.0)
|
||||
parser.add_argument("--count", type=int, default=10)
|
||||
parser.add_argument("--interval", type=float, default=1.0)
|
||||
parser.add_argument("--old-token", type=int)
|
||||
parser.add_argument("--old-upf", type=int)
|
||||
args = parser.parse_args()
|
||||
if args.timeout <= 0 or args.count < 0 or args.interval < 0:
|
||||
parser.error("timeout must be positive; count and interval must be nonnegative")
|
||||
if (args.old_token is None) != (args.old_upf is None):
|
||||
parser.error("--old-token and --old-upf must be supplied together")
|
||||
try:
|
||||
asyncio.run(run(args))
|
||||
except (OSError, TimeoutError, ProtocolError) as error:
|
||||
parser.exit(1, f"Client failed: {error}\n")
|
||||
except KeyboardInterrupt:
|
||||
parser.exit(130, "Client stopped.\n")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
10
examples/start_client.bat
Normal file
10
examples/start_client.bat
Normal file
@@ -0,0 +1,10 @@
|
||||
@echo off
|
||||
setlocal
|
||||
cd /d "%~dp0.."
|
||||
set "PYTHON=D:\_Programs\anaconda\envs\sites\python.exe"
|
||||
if not exist "%PYTHON%" (
|
||||
echo Python not found: %PYTHON%
|
||||
exit /b 1
|
||||
)
|
||||
"%PYTHON%" -u "%~dp0client.py" %*
|
||||
exit /b %ERRORLEVEL%
|
||||
Reference in New Issue
Block a user