[+!] all but small

some small modifications, tests added

examples/client.py

@@ -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()