progpib/tests/conftest.py

import socket
import threading

import pytest


class MockTcpServer:
    """A single-connection TCP server on 127.0.0.1, driven by `handler(conn)`
    in a background thread. Used to test transports/controller without
    hardware."""

    def __init__(self, handler):
        self._handler = handler
        self._listen_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self._listen_sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        self._listen_sock.bind(("127.0.0.1", 0))
        self._listen_sock.listen(1)
        self.host, self.port = self._listen_sock.getsockname()
        self._thread = threading.Thread(target=self._serve, daemon=True)
        self._thread.start()

    def _serve(self) -> None:
        try:
            conn, _ = self._listen_sock.accept()
        except OSError:
            return
        try:
            self._handler(conn)
        finally:
            conn.close()

    def close(self) -> None:
        self._listen_sock.close()
        self._thread.join(timeout=2.0)


@pytest.fixture
def mock_server():
    servers: list[MockTcpServer] = []

    def make(handler) -> MockTcpServer:
        server = MockTcpServer(handler)
        servers.append(server)
        return server

    yield make
    for server in servers:
        server.close()


class FakeGpibServer:
    """Emulates enough of the Prologix "++" command protocol to unit-test
    `GpibController` without hardware.

    Tracks config command state in `self.state` and answers query/set
    semantics per the Prologix manual. Also records every raw byte
    received in `self.raw_received`, independent of line parsing, so tests
    can check exact wire bytes (e.g. escaping) without the naive
    line-splitter tripping over an escaped LF/CR inside a payload.
    """

    def __init__(self, *, version: str = "Fake Prologix GPIB-ETHERNET Version 1.0"):
        self.state = {
            "mode": "1",
            "addr": "0",
            "auto": "0",
            "eoi": "1",
            "eos": "0",
            "eot_enable": "0",
            "eot_char": "0",
            "read_tmo_ms": "500",
            "savecfg": "1",
            "status": "0",
            "lon": "0",
        }
        self.version = version
        self.spoll_response = 0
        self.read_response = b""
        self.raw_received = bytearray()

    def handle(self, conn: socket.socket) -> None:
        buf = bytearray()
        conn.settimeout(5.0)
        while True:
            try:
                chunk = conn.recv(4096)
            except (socket.timeout, OSError):
                break
            if not chunk:
                break
            self.raw_received.extend(chunk)
            buf.extend(chunk)
            while b"\n" in buf:
                idx = buf.index(b"\n")
                line = bytes(buf[:idx])
                del buf[: idx + 1]
                self._handle_line(conn, line)

    def _handle_line(self, conn: socket.socket, line: bytes) -> None:
        if not line.startswith(b"++"):
            return  # raw instrument data; use raw_received to inspect it
        self._handle_command(conn, line[2:].decode("ascii"))

    def _handle_command(self, conn: socket.socket, command: str) -> None:
        parts = command.split()
        if not parts:
            return
        name, args = parts[0], parts[1:]

        if name == "ver":
            conn.sendall((self.version + "\n").encode("ascii"))
        elif name == "read":
            conn.sendall(self.read_response)
        elif name == "spoll":
            conn.sendall((str(self.spoll_response) + "\n").encode("ascii"))
        elif name == "addr":
            if args:
                self.state["addr"] = " ".join(args)
            else:
                conn.sendall((self.state["addr"] + "\n").encode("ascii"))
        elif name in self.state:
            if args:
                self.state[name] = args[0]
            else:
                conn.sendall((str(self.state[name]) + "\n").encode("ascii"))
        # else: silently ignore action commands with no return value we
        # don't need for these tests (clr, ifc, llo, loc, rst, srq, trg, help)


@pytest.fixture
def fake_gpib_server(mock_server):
    server_obj = FakeGpibServer()
    tcp_server = mock_server(server_obj.handle)
    return server_obj, tcp_server