diff --git a/src/comet/emulator/__main__.py b/src/comet/emulator/__main__.py index 4f6f2d2..907a88c 100644 --- a/src/comet/emulator/__main__.py +++ b/src/comet/emulator/__main__.py @@ -11,10 +11,6 @@ lcr: model: urn:comet:model:keysight:e4980a port: 11002 - # User specific emulator - my_instr: - module: local_project.my_instr_emulator - port: 12001 ``` Loading a configuration filename (default filenames are `emulators.yaml` and `emulators.yml`). @@ -28,19 +24,19 @@ """ import argparse +import asyncio +import contextlib import logging import os import signal -import threading from typing import Any import schema import yaml from .. import __version__ - from .emulator import emulator_factory -from .tcpserver import TCPServer, TCPServerThread, TCPServerContext +from .tcpserver import TCPServer, TCPServerContext default_config_filenames: list[str] = ["emulators.yaml", "emulators.yml"] default_host: str = "localhost" @@ -68,7 +64,7 @@ def normalize_termination(value: str) -> str: config_schema = schema.Schema( { - schema.Optional("version"): str, # deprecated + schema.Optional("version"): str, "emulators": { str: { schema.Optional(schema.Or("model", "module")): str, @@ -93,10 +89,15 @@ def load_config(filename: str) -> dict[str, Any]: with open(filename) as fp: data = yaml.safe_load(fp) config = validate_config(data or {}) - # Set defaults - for params in config.get("emulators", {}).values(): + for name, params in config.get("emulators", {}).items(): if "model" in params and "module" in params: raise KeyError("keys 'model' and 'module' are exclusive") + if "module" in params: + logging.warning( + "Emulator %r uses deprecated config key 'module'; " + "use 'model' instead. Support exists only for backward compatibility.", + name, + ) params.setdefault("host", default_host) params.setdefault("termination", default_termination) params.setdefault("request_delay", default_request_delay) @@ -132,37 +133,26 @@ def locate_config_filename() -> str: raise RuntimeError("No config file found.") -def event_loop() -> None: - """Blocks execution until termination or interrupt from keyboard signal.""" - e = threading.Event() - - def handle_event(signum, frame): - e.set() - - signal.signal(signal.SIGTERM, handle_event) - signal.signal(signal.SIGINT, handle_event) - e.wait() - - -def main() -> None: +async def main() -> None: args = parse_args() logging.basicConfig(level=logging.INFO) config = load_config(args.filename or locate_config_filename()) - threads = [] + servers: list[TCPServer] = [] for name, params in config.get("emulators", {}).items(): - model = params.get("model") or params.get("module") # fallback for comet<1.5 + model = params.get("model") or params.get("module") # fallback for comet<1.5 host = params.get("host") port = params.get("port") termination_bytes = params.get("termination").encode() request_delay = params.get("request_delay") options = params.get("options", {}) - address = host, port + emulator = emulator_factory(model)() emulator.options.update(options) + context = TCPServerContext( name=name, emulator=emulator, @@ -170,26 +160,51 @@ def main() -> None: request_delay=request_delay, logger=logging.getLogger(name), ) - server = TCPServer(address, context) - threads.append(TCPServerThread(server)) - - for thread in threads: - host, port, *_ = thread.server.server_address # IPv4/IPv6 - thread.server.context.logger.info("starting... %s:%s", host, port) - thread.start() - - def handle_event(signum, frame): - for thread in threads: - host, port, *_ = thread.server.server_address # IPv4/IPv6 - thread.server.context.logger.info("stopping... %s:%s", host, port) - thread.shutdown() - - signal.signal(signal.SIGTERM, handle_event) - signal.signal(signal.SIGINT, handle_event) - - for thread in threads: - thread.join() + server = TCPServer((host, port), context) + await server.start() + servers.append(server) + + for server in servers: + host, port = server.server_address + server.context.logger.info("starting... %s:%s", host, port) + + stop_event = asyncio.Event() + + def request_shutdown() -> None: + if not stop_event.is_set(): + for server in servers: + host, port = server.server_address + server.context.logger.info("stopping... %s:%s", host, port) + stop_event.set() + + loop = asyncio.get_running_loop() + + for sig in (signal.SIGTERM, signal.SIGINT): + try: + loop.add_signal_handler(sig, request_shutdown) + except NotImplementedError: + # Not supported on some platforms (notably parts of Windows). + # In that case, asyncio.run() will still surface Ctrl+C as + # KeyboardInterrupt, handled outside main(). + ... + + tasks = [asyncio.create_task(server.serve_forever()) for server in servers] + + try: + await stop_event.wait() + finally: + for server in servers: + await server.shutdown() + for task in tasks: + task.cancel() + for task in tasks: + with contextlib.suppress(asyncio.CancelledError): + await task if __name__ == "__main__": - main() + try: + asyncio.run(main()) + except KeyboardInterrupt: + # Fallback for platforms where asyncio signal handlers are unavailable. + ... diff --git a/src/comet/emulator/tcpserver.py b/src/comet/emulator/tcpserver.py index a59d1f4..2dc4068 100644 --- a/src/comet/emulator/tcpserver.py +++ b/src/comet/emulator/tcpserver.py @@ -1,10 +1,9 @@ import argparse +import asyncio +import contextlib import inspect import logging import re -import socketserver -import threading -import time import signal from dataclasses import dataclass from typing import Iterable, Optional, Union @@ -12,48 +11,81 @@ from .emulator import Emulator from .response import Response -__all__ = ["TCPRequestHandler", "TCPServer", "TCPServerThread", "TCPServerContext"] +__all__ = ["TCPRequestHandler", "TCPServer", "TCPServerContext"] -class TCPRequestHandler(socketserver.BaseRequestHandler): - def read_messages(self, context: "TCPServerContext", rx_buffer: bytearray) -> Optional[list[bytes]]: - data = self.request.recv(4096) +class TCPRequestHandler: + async def read_messages( + self, + reader: asyncio.StreamReader, + context: "TCPServerContext", + rx_buffer: bytearray, + ) -> Optional[list[bytes]]: + data = await reader.read(4096) if not data: return None + context.logger.info("recv %s", data) rx_buffer.extend(data) + termination_bytes = context.termination messages: list[bytes] = [] + while (pos := rx_buffer.find(termination_bytes)) >= 0: message = rx_buffer[:pos] - del rx_buffer[:pos + len(termination_bytes)] + del rx_buffer[: pos + len(termination_bytes)] messages.append(bytes(message)) + return messages - def send_messages(self, context: "TCPServerContext", response: Union[Response, Iterable[Response]]) -> None: + async def send_messages( + self, + writer: asyncio.StreamWriter, + context: "TCPServerContext", + response: Union[Response, Iterable[Response]], + ) -> None: termination_bytes = context.termination - if not isinstance(response, (list, tuple)): - response = [response] # type: ignore + + if isinstance(response, Response): + responses = [response] + else: + responses = list(response) + buffer = bytearray() - for res in response: + for res in responses: buffer.extend(bytes(res)) - # TODO most instruments append termination also to binary + # TODO: most instruments append termination also to binary buffer.extend(termination_bytes) + data = bytes(buffer) context.logger.info("send %s", data) - self.request.sendall(data) - - def handle(self) -> None: - context = self.server.context # type: ignore + writer.write(data) + await writer.drain() + + async def handle( + self, + reader: asyncio.StreamReader, + writer: asyncio.StreamWriter, + context: "TCPServerContext", + ) -> None: rx_buffer = bytearray() - while True: - messages = self.read_messages(context, rx_buffer) - if messages is None: - break - for message in messages: - response = context.handle_message(str(message, "utf-8")) - if response is not None: - self.send_messages(context, response) + + try: + while True: + messages = await self.read_messages(reader, context, rx_buffer) + if messages is None: + break + + for message in messages: + response = await context.handle_message(str(message, "utf-8")) + if response is not None: + await self.send_messages(writer, context, response) + except asyncio.CancelledError: + raise + finally: + writer.close() + with contextlib.suppress(Exception): + await writer.wait_closed() @dataclass @@ -64,56 +96,75 @@ class TCPServerContext: request_delay: float logger: logging.Logger - def handle_message(self, message: str) -> Union[None, Response, Iterable[Response]]: + async def handle_message( + self, + message: str, + ) -> Union[None, Response, Iterable[Response]]: response = self.emulator(message) - time.sleep(self.request_delay) + if response is not None: + await asyncio.sleep(self.request_delay) return response -class TCPServer(socketserver.TCPServer): - allow_reuse_address: bool = True - +class TCPServer: def __init__(self, address: tuple[str, int], context: TCPServerContext) -> None: - super().__init__(address, TCPRequestHandler) - self.context: TCPServerContext = context - - -class TCPServerThread(threading.Thread): - def __init__(self, server: TCPServer) -> None: - super().__init__() - self.server: TCPServer = server - self._shutdown_request: threading.Event = threading.Event() - - def shutdown(self) -> None: - self._shutdown_request.set() - self.server.shutdown() - - def run(self) -> None: - with self.server as server: - server.serve_forever() + self.address = address + self.context = context + self._server: Optional[asyncio.base_events.Server] = None + self._handler = TCPRequestHandler() + self._shutdown_lock = asyncio.Lock() + self._shutdown_started = False + + @property + def server_address(self) -> tuple[str, int]: + if self._server is None or not self._server.sockets: + return self.address + + sockname = self._server.sockets[0].getsockname() + return sockname[:2] + + async def start(self) -> None: + if self._server is not None: + return + + self._server = await asyncio.start_server( + lambda r, w: self._handler.handle(r, w, self.context), + host=self.address[0], + port=self.address[1], + reuse_address=True, + ) + + async def serve_forever(self) -> None: + if self._server is None: + await self.start() + + assert self._server is not None + async with self._server: + await self._server.serve_forever() + + async def shutdown(self) -> None: + async with self._shutdown_lock: + if self._shutdown_started: + return + self._shutdown_started = True + + if self._server is not None: + self._server.close() + await self._server.wait_closed() def option_type(value: str) -> tuple[str, str]: - m = re.match(r"^([\w_][\w\d_]*)=(.*)$", value) + m = re.match(r"^([A-Za-z0-9][A-Za-z0-9_.:+/-]*)=(.*)$", value) if m: return m.group(1), m.group(2) - raise argparse.ArgumentTypeError("expected key=value") + + raise argparse.ArgumentTypeError("expected 'key=value'") def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser() - parser.add_argument( - "--host", - default="localhost", - help="host, default is 'localhost'", - ) - parser.add_argument( - "-p", - "--port", - type=int, - default=10000, - help="port, default is 10000", - ) + parser.add_argument("--host", default="localhost", help="host, default is 'localhost'") + parser.add_argument("-p", "--port", type=int, default=10000, help="port, default is 10000") parser.add_argument( "-t", "--termination", @@ -140,7 +191,7 @@ def parse_args() -> argparse.Namespace: def run(emulator: Emulator) -> int: - """Convenience emulator runner using TCP server.""" + """Convenience emulator runner using asyncio TCP server.""" if not isinstance(emulator, Emulator): raise TypeError(f"Emulator must inherit from {Emulator}") @@ -150,8 +201,10 @@ def run(emulator: Emulator) -> int: logging.basicConfig(level=logging.INFO) mod = inspect.getmodule(emulator.__class__) - name = (getattr(getattr(mod, "__spec__", None), "name", None) - or emulator.__class__.__module__) + name = ( + getattr(getattr(mod, "__spec__", None), "name", None) + or emulator.__class__.__module__ + ) context = TCPServerContext( name=name, @@ -160,22 +213,52 @@ def run(emulator: Emulator) -> int: request_delay=args.request_delay, logger=logging.getLogger(name), ) - address = args.host, args.port - server = TCPServer(address, context) - thread = TCPServerThread(server) + address = (args.host, args.port) + + async def main() -> None: + server = TCPServer(address, context) + await server.start() + + host, port = server.server_address + context.logger.info("starting... %s:%s", host, port) + + stop_event = asyncio.Event() + + def request_shutdown() -> None: + if stop_event.is_set(): + return + + context.logger.info("stopping... %s:%s", host, port) + stop_event.set() + + loop = asyncio.get_running_loop() + installed_signals: list[int] = [] - host, port, *_ = server.server_address # IPv4/IPv6 - context.logger.info("starting... %s:%s", host, port) + for sig in (signal.SIGTERM, signal.SIGINT): + try: + loop.add_signal_handler(sig, request_shutdown) + installed_signals.append(sig) + except NotImplementedError: + # Not supported on some platforms (notably parts of Windows). + # In that case, asyncio.run() will still surface Ctrl+C as + # KeyboardInterrupt, handled outside main(). + ... - thread.start() + serve_task = asyncio.create_task(server.serve_forever()) - def handle_event(signum, frame): - context.logger.info("stopping... %s:%s", host, port) - server.shutdown() + try: + await stop_event.wait() + finally: + await server.shutdown() - signal.signal(signal.SIGTERM, handle_event) - signal.signal(signal.SIGINT, handle_event) + serve_task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await serve_task - thread.join() + try: + asyncio.run(main()) + except KeyboardInterrupt: + # Fallback for platforms where asyncio signal handlers are unavailable. + ... return 0 diff --git a/tests/test_emulator_tcpserver.py b/tests/test_emulator_tcpserver.py new file mode 100644 index 0000000..253d824 --- /dev/null +++ b/tests/test_emulator_tcpserver.py @@ -0,0 +1,386 @@ +import argparse +import logging +import types + +import pytest + +from comet.emulator import tcpserver + + +class FakeResponse: + def __init__(self, payload: bytes): + self.payload = payload + + def __bytes__(self) -> bytes: + return self.payload + + +class FakeStreamReader: + def __init__(self, chunks): + self._chunks = list(chunks) + + async def read(self, n: int) -> bytes: + if self._chunks: + return self._chunks.pop(0) + return b"" + + +class FakeStreamWriter: + def __init__(self): + self.writes = [] + self.drained = 0 + self.closed = False + self.wait_closed_called = False + + def write(self, data: bytes) -> None: + self.writes.append(data) + + async def drain(self) -> None: + self.drained += 1 + + def close(self) -> None: + self.closed = True + + async def wait_closed(self) -> None: + self.wait_closed_called = True + + +class DummyLogger: + def __init__(self): + self.records = [] + + def info(self, msg, *args): + if args: + msg = msg % args + self.records.append(msg) + + +class DummyContext: + def __init__(self, termination=b"\n", handle_result=None): + self.termination = termination + self.logger = DummyLogger() + self._handle_result = handle_result + self.messages = [] + + async def handle_message(self, message: str): + self.messages.append(message) + return self._handle_result + + +@pytest.fixture +def patch_response(monkeypatch): + monkeypatch.setattr(tcpserver, "Response", FakeResponse) + return FakeResponse + + +@pytest.mark.asyncio +async def test_read_messages_returns_none_on_eof(): + handler = tcpserver.TCPRequestHandler() + reader = FakeStreamReader([b""]) + context = DummyContext() + rx_buffer = bytearray() + + result = await handler.read_messages(reader, context, rx_buffer) + + assert result is None + assert rx_buffer == bytearray() + assert context.logger.records == [] + + +@pytest.mark.asyncio +async def test_read_messages_splits_complete_messages(): + handler = tcpserver.TCPRequestHandler() + reader = FakeStreamReader([b"CMD1\nCMD2\n"]) + context = DummyContext(termination=b"\n") + rx_buffer = bytearray() + + result = await handler.read_messages(reader, context, rx_buffer) + + assert result == [b"CMD1", b"CMD2"] + assert rx_buffer == bytearray() + assert context.logger.records == ["recv b'CMD1\\nCMD2\\n'"] + + +@pytest.mark.asyncio +async def test_read_messages_preserves_partial_frame(): + handler = tcpserver.TCPRequestHandler() + context = DummyContext(termination=b"\r\n") + rx_buffer = bytearray() + + result1 = await handler.read_messages( + FakeStreamReader([b"ONE\r\nTWO\r"]), + context, + rx_buffer, + ) + assert result1 == [b"ONE"] + assert rx_buffer == bytearray(b"TWO\r") + + result2 = await handler.read_messages( + FakeStreamReader([b"\nTHREE\r\n"]), + context, + rx_buffer, + ) + assert result2 == [b"TWO", b"THREE"] + assert rx_buffer == bytearray() + + +@pytest.mark.asyncio +async def test_send_messages_single_response(patch_response): + handler = tcpserver.TCPRequestHandler() + writer = FakeStreamWriter() + context = DummyContext(termination=b"\n") + + await handler.send_messages(writer, context, FakeResponse(b"OK")) + + assert writer.writes == [b"OK\n"] + assert writer.drained == 1 + assert context.logger.records == ["send b'OK\\n'"] + + +@pytest.mark.asyncio +async def test_send_messages_multiple_responses(patch_response): + handler = tcpserver.TCPRequestHandler() + writer = FakeStreamWriter() + context = DummyContext(termination=b"\r\n") + + responses = [FakeResponse(b"A"), FakeResponse(b"B")] + await handler.send_messages(writer, context, responses) + + assert writer.writes == [b"A\r\nB\r\n"] + assert writer.drained == 1 + assert context.logger.records == ["send b'A\\r\\nB\\r\\n'"] + + +@pytest.mark.asyncio +async def test_handle_reads_dispatches_sends_and_closes_writer(monkeypatch, patch_response): + handler = tcpserver.TCPRequestHandler() + context = DummyContext(handle_result=FakeResponse(b"RSP")) + writer = FakeStreamWriter() + + async def fake_read_messages(reader, context, rx_buffer): + if not hasattr(fake_read_messages, "count"): + fake_read_messages.count = 0 + fake_read_messages.count += 1 + if fake_read_messages.count == 1: + return [b"PING", b"PONG"] + return None + + sent = [] + + async def fake_send_messages(writer, context, response): + sent.append(bytes(response)) + + monkeypatch.setattr(handler, "read_messages", fake_read_messages) + monkeypatch.setattr(handler, "send_messages", fake_send_messages) + + await handler.handle(object(), writer, context) + + assert context.messages == ["PING", "PONG"] + assert sent == [b"RSP", b"RSP"] + assert writer.closed is True + assert writer.wait_closed_called is True + + +@pytest.mark.asyncio +async def test_handle_closes_writer_even_if_handle_message_raises(): + handler = tcpserver.TCPRequestHandler() + writer = FakeStreamWriter() + + class ExplodingContext(DummyContext): + async def handle_message(self, message: str): + raise RuntimeError("boom") + + context = ExplodingContext() + + async def fake_read_messages(reader, context, rx_buffer): + return [b"PING"] + + handler.read_messages = fake_read_messages + + with pytest.raises(RuntimeError, match="boom"): + await handler.handle(object(), writer, context) + + assert writer.closed is True + assert writer.wait_closed_called is True + + +@pytest.mark.asyncio +async def test_context_handle_message_waits_only_when_response(monkeypatch): + sleep_calls = [] + + async def fake_sleep(delay): + sleep_calls.append(delay) + + monkeypatch.setattr(tcpserver.asyncio, "sleep", fake_sleep) + + class DummyEmulator: + def __init__(self, response): + self.response = response + + def __call__(self, message): + return self.response + + ctx1 = tcpserver.TCPServerContext( + name="x", + emulator=DummyEmulator("ok"), + termination=b"\n", + request_delay=0.25, + logger=logging.getLogger("test1"), + ) + ctx2 = tcpserver.TCPServerContext( + name="x", + emulator=DummyEmulator(None), + termination=b"\n", + request_delay=0.25, + logger=logging.getLogger("test2"), + ) + + assert await ctx1.handle_message("cmd") == "ok" + assert sleep_calls == [0.25] + + assert await ctx2.handle_message("cmd") is None + assert sleep_calls == [0.25] + + +def test_server_address_returns_configured_address_before_start(): + ctx = DummyContext() + server = tcpserver.TCPServer(("127.0.0.1", 5555), ctx) + + assert server.server_address == ("127.0.0.1", 5555) + + +def test_server_address_returns_bound_socket_address_after_start(): + ctx = DummyContext() + server = tcpserver.TCPServer(("127.0.0.1", 5555), ctx) + + class DummySocket: + def getsockname(self): + return ("127.0.0.1", 6000, "ignored") + + server._server = types.SimpleNamespace(sockets=[DummySocket()]) + + assert server.server_address == ("127.0.0.1", 6000) + + +@pytest.mark.asyncio +async def test_server_start_uses_asyncio_start_server(monkeypatch): + ctx = DummyContext() + server = tcpserver.TCPServer(("localhost", 1234), ctx) + + created = {} + + class DummyAsyncServer: + sockets = [] + + async def fake_start_server(callback, host, port, reuse_address): + created["callback"] = callback + created["host"] = host + created["port"] = port + created["reuse_address"] = reuse_address + return DummyAsyncServer() + + monkeypatch.setattr(tcpserver.asyncio, "start_server", fake_start_server) + + await server.start() + + assert isinstance(server._server, DummyAsyncServer) + assert created["host"] == "localhost" + assert created["port"] == 1234 + assert created["reuse_address"] is True + assert callable(created["callback"]) + + +@pytest.mark.asyncio +async def test_server_shutdown_closes_server(): + ctx = DummyContext() + server = tcpserver.TCPServer(("localhost", 1234), ctx) + + class DummyAsyncServer: + def __init__(self): + self.closed = False + self.wait_closed_called = False + + def close(self): + self.closed = True + + async def wait_closed(self): + self.wait_closed_called = True + + dummy = DummyAsyncServer() + server._server = dummy + + await server.shutdown() + + assert dummy.closed is True + assert dummy.wait_closed_called is True + + +def test_option_type_parses_valid_key_value(): + assert tcpserver.option_type("version=2.1") == ("version", "2.1") + assert tcpserver.option_type("a_b.c:+/-=value") == ("a_b.c:+/-", "value") + + +def test_option_type_rejects_invalid_key(): + with pytest.raises(argparse.ArgumentTypeError, match="expected 'key=value'"): + tcpserver.option_type("=value") + + +def test_parse_args_defaults(monkeypatch): + monkeypatch.setattr(tcpserver.argparse.ArgumentParser, "parse_args", lambda self: argparse.Namespace( + host="localhost", + port=10000, + termination="\n", + request_delay=0.1, + option=[], + )) + + args = tcpserver.parse_args() + + assert args.host == "localhost" + assert args.port == 10000 + assert args.termination == "\n" + assert args.request_delay == 0.1 + assert args.option == [] + + +def test_run_rejects_non_emulator(): + with pytest.raises(TypeError, match="Emulator must inherit from"): + tcpserver.run(object()) + + +def test_run_configures_options_and_returns_zero(monkeypatch): + class BaseEmulator: + def __init__(self): + self.options = {} + + monkeypatch.setattr(tcpserver, "Emulator", BaseEmulator) + + emulator = BaseEmulator() + + monkeypatch.setattr( + tcpserver, + "parse_args", + lambda: argparse.Namespace( + host="127.0.0.1", + port=9999, + termination="\n", + request_delay=0.01, + option=[("mode", "test"), ("version", "1")], + ), + ) + + monkeypatch.setattr(tcpserver.inspect, "getmodule", lambda cls: types.SimpleNamespace(__spec__=types.SimpleNamespace(name="pkg.tcpserver"))) + + ran = {} + + def fake_asyncio_run(coro): + ran["called"] = True + coro.close() + + monkeypatch.setattr(tcpserver.asyncio, "run", fake_asyncio_run) + + rc = tcpserver.run(emulator) + + assert rc == 0 + assert ran["called"] is True + assert emulator.options == {"mode": "test", "version": "1"} diff --git a/tox.ini b/tox.ini index 084af0e..26d8ee5 100644 --- a/tox.ini +++ b/tox.ini @@ -9,6 +9,7 @@ deps = mypy types-PyYAML pytest + pytest-asyncio commands = ruff check src tests mypy src