esphome/tests/integration/raw_api_client.py

"""Minimal plaintext native-api client over a raw socket.

Reads only when told to, so tests control when the TCP pipe backs up toward
the device. Message types are counted, and payloads are kept for read_frame().
"""

from __future__ import annotations

import asyncio
from collections import Counter, deque
import socket
from typing import Self

from aioesphomeapi import api_pb2
import aioesphomeapi.core as api_core
from google.protobuf import message

from .const import LOCALHOST

# Message type ids are protocol constants; derive them from aioesphomeapi so
# they cannot drift from the client library in use.
MESSAGE_TYPE_OF = {cls: num for num, cls in api_core.MESSAGE_TYPE_TO_PROTO.items()}

_READ_CHUNK = 4096


def encode_varint(value: int) -> bytes:
    out = bytearray()
    while True:
        byte = value & 0x7F
        value >>= 7
        if value:
            out.append(byte | 0x80)
        else:
            out.append(byte)
            return bytes(out)


def decode_varint(buf: bytearray, pos: int) -> tuple[int, int] | None:
    """Decode one varint at pos; return (value, new_pos) or None if short."""
    value = shift = 0
    while pos < len(buf):
        byte = buf[pos]
        pos += 1
        value |= (byte & 0x7F) << shift
        if not byte & 0x80:
            return value, pos
        shift += 7
    return None


def decode_fields(payload: bytes) -> dict[int, list[int | bytes]]:
    """Split a protobuf payload into its field values by field number.

    For messages the installed aioesphomeapi does not know yet.
    """
    fields: dict[int, list[int | bytes]] = {}
    buf = bytearray(payload)

    def varint(pos: int) -> tuple[int, int]:
        decoded = decode_varint(buf, pos)
        assert decoded is not None, "truncated varint in payload"
        return decoded

    pos = 0
    while pos < len(buf):
        tag, pos = varint(pos)
        wire_type = tag & 0x07
        value: int | bytes
        if wire_type == 0:
            value, pos = varint(pos)
        elif wire_type == 2:
            length, pos = varint(pos)
            value, pos = bytes(buf[pos : pos + length]), pos + length
        elif wire_type == 5:
            value, pos = int.from_bytes(buf[pos : pos + 4], "little"), pos + 4
        else:
            raise AssertionError(f"unexpected wire type {wire_type}")
        fields.setdefault(tag >> 3, []).append(value)
    return fields


def encode_frame(msg_type: int, payload: bytes) -> bytes:
    """Encode one plaintext api frame: 0x00, payload length, message type."""
    return b"\x00" + encode_varint(len(payload)) + encode_varint(msg_type) + payload


class FrameParser:
    """Incremental parser for the plaintext api frame stream."""

    def __init__(self) -> None:
        self._buf = bytearray()
        self.frames: deque[tuple[int, bytes]] = deque()

    def feed(self, data: bytes) -> list[int]:
        self._buf.extend(data)
        types: list[int] = []
        while (frame := self._try_parse()) is not None:
            self.frames.append(frame)
            types.append(frame[0])
        return types

    def _try_parse(self) -> tuple[int, bytes] | None:
        buf = self._buf
        if not buf:
            return None
        assert buf[0] == 0, f"expected plaintext frame, got indicator {buf[0]}"
        if (size_decoded := decode_varint(buf, 1)) is None:
            return None
        size, pos = size_decoded
        if (type_decoded := decode_varint(buf, pos)) is None:
            return None
        msg_type, pos = type_decoded
        if len(buf) - pos < size:
            return None
        payload = bytes(buf[pos : pos + size])
        del buf[: pos + size]
        return msg_type, payload


class RawApiClient:
    """Plaintext api client whose reads happen only on request."""

    def __init__(self, port: int, recv_buffer_size: int | None = None) -> None:
        self._port = port
        self._parser = FrameParser()
        self.bytes_received = 0
        self.frame_counts: Counter[int] = Counter()
        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        try:
            if recv_buffer_size is not None:
                sock.setsockopt(socket.SOL_SOCKET, socket.SO_RCVBUF, recv_buffer_size)
                # Kernels may round up (Linux doubles) but must not clamp below
                applied = sock.getsockopt(socket.SOL_SOCKET, socket.SO_RCVBUF)
                assert applied >= recv_buffer_size, (
                    f"SO_RCVBUF clamped to {applied}, requested {recv_buffer_size}"
                )
            sock.setblocking(False)
        except Exception:
            sock.close()
            raise
        self._sock = sock

    async def __aenter__(self) -> Self:
        return self

    async def __aexit__(self, *exc_info: object) -> None:
        self.close()

    async def connect(self, client_info: str = "raw-api-client") -> None:
        """Connect and complete the Hello handshake (no auth step since 2026.1.0)."""
        loop = asyncio.get_running_loop()
        await loop.sock_connect(self._sock, (LOCALHOST, self._port))
        hello = api_pb2.HelloRequest()
        hello.client_info = client_info
        hello.api_version_major = 1
        hello.api_version_minor = 10
        await self.send_message(hello)
        await self.read_until_frame(MESSAGE_TYPE_OF[api_pb2.HelloResponse])

    async def send_message(self, msg: message.Message) -> None:
        await self.send_raw(MESSAGE_TYPE_OF[type(msg)], msg.SerializeToString())

    async def send_raw(self, msg_type: int, payload: bytes) -> None:
        """Send a frame with a hand built payload, for shapes protobuf will not serialize."""
        loop = asyncio.get_running_loop()
        await loop.sock_sendall(self._sock, encode_frame(msg_type, payload))

    async def read_until_frame(self, msg_type: int, timeout: float = 10.0) -> None:
        """Read until at least one frame of msg_type has been received."""
        loop = asyncio.get_running_loop()

        async def _read_loop() -> None:
            while not self.frame_counts[msg_type]:
                data = await loop.sock_recv(self._sock, _READ_CHUNK)
                assert data, "server closed the connection unexpectedly"
                self.bytes_received += len(data)
                self.frame_counts.update(self._parser.feed(data))

        await asyncio.wait_for(_read_loop(), timeout)

    async def read_frame(self, msg_type: int, timeout: float = 10.0) -> bytes:
        """Return the payload of the next frame of msg_type, dropping other frames before it."""
        loop = asyncio.get_running_loop()
        frames = self._parser.frames

        async def _read_loop() -> bytes:
            while True:
                while frames:
                    frame_type, payload = frames.popleft()
                    if frame_type == msg_type:
                        return payload
                data = await loop.sock_recv(self._sock, _READ_CHUNK)
                assert data, "server closed the connection unexpectedly"
                self.bytes_received += len(data)
                self.frame_counts.update(self._parser.feed(data))

        return await asyncio.wait_for(_read_loop(), timeout)

    def close(self) -> None:
        self._sock.close()