core/examples/client.py

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