dynalab-core/test/test_json_server.py

import asyncio
import logging
import socket
import time
from uuid import uuid4

import pytest
from dynalab_core import Core
from dynalab_core.config import CoreConfig
from dynalab_core.errors import JsonServerError
from dynalab_core.protocols.constants import PROTOCOL_VERSION
from dynalab_core.protocols.common import VersionDescriptor
from dynalab_core.protocols.packets.handshake import ConnectorHello, SignalDescriptor
from test.common import find_available_port


def _wait_for_event(
    caplog: pytest.LogCaptureFixture, event: str, timeout: float = 1.0
) -> logging.LogRecord:
    deadline = time.monotonic() + timeout
    while time.monotonic() < deadline:
        for record in caplog.records:
            if getattr(record, "event", None) == event:
                return record
        time.sleep(0.01)
    raise AssertionError(f"Log event {event!r} was not emitted")


def test_json_server_raises_timeout_error(
    caplog: pytest.LogCaptureFixture,
) -> None:
    port = find_available_port(8765)
    core = Core(CoreConfig(port=port))
    core._json_server._debug_timeout_test = True
    core._json_server._timeout = 0.01

    with caplog.at_level(logging.DEBUG, logger="dynalab_core"):
        try:
            with pytest.raises(JsonServerError, match="did not start") as raised:
                core.start()
        finally:
            core.stop()

    timeout_record = next(
        record
        for record in caplog.records
        if getattr(record, "event", None) == "json_server.start_timeout"
    )
    assert timeout_record.levelno == logging.ERROR
    assert timeout_record.port == port
    assert raised.value.code == "start_timeout"


def test_json_server_raises_startup_error(
    caplog: pytest.LogCaptureFixture,
) -> None:

    port = find_available_port(8765)
    core1 = Core(CoreConfig(port=port))
    core2 = Core(CoreConfig(port=port))

    with caplog.at_level(logging.DEBUG, logger="dynalab_core"):
        core1.start()

        try:
            with pytest.raises(JsonServerError, match="failed to start") as raised:
                core2.start()
        finally:
            core1.stop()
            core2.stop()

    failure_record = next(
        record
        for record in caplog.records
        if getattr(record, "event", None) == "json_server.start_failed"
    )
    assert failure_record.levelno == logging.ERROR
    assert failure_record.port == port
    assert failure_record.exc_info is not None
    assert raised.value.code == "startup_failed"
    assert isinstance(raised.value.__cause__, OSError)


def test_json_server_logs_invalid_handshake(
    caplog: pytest.LogCaptureFixture,
) -> None:
    port = find_available_port(8765)
    core = Core(CoreConfig(port=port))

    with caplog.at_level(logging.DEBUG, logger="dynalab_core"):
        core.start()
        try:
            with socket.create_connection(("127.0.0.1", port), timeout=1) as peer:
                server_hello = peer.makefile("rb").readline()
                assert server_hello
                peer.sendall(b'{"secret":"must-not-appear-in-logs"}\n')

                rejection = _wait_for_event(caplog, "connector.handshake_rejected")
        finally:
            core.stop()

    assert rejection.levelno == logging.WARNING
    assert rejection.reason == "invalid_frame"
    assert rejection.connection_id
    assert rejection.peer_address
    assert "must-not-appear-in-logs" not in caplog.text


def test_json_server_reports_unexpected_thread_startup_failure(
    caplog: pytest.LogCaptureFixture, monkeypatch: pytest.MonkeyPatch
) -> None:
    async def fail_start_server(*args: object, **kwargs: object) -> None:
        raise RuntimeError("unexpected startup failure")

    monkeypatch.setattr(asyncio, "start_server", fail_start_server)
    core = Core(CoreConfig(port=find_available_port(8765)))

    with caplog.at_level(logging.DEBUG, logger="dynalab_core"):
        try:
            with pytest.raises(JsonServerError) as raised:
                core.start()
        finally:
            core.stop()

    failures = [
        record
        for record in caplog.records
        if getattr(record, "event", None) == "json_server.thread_failed"
    ]
    assert len(failures) == 1
    assert failures[0].startup_complete is False
    assert failures[0].exception_type == "RuntimeError"
    assert failures[0].exc_info is not None
    assert raised.value.code == "startup_failed"
    assert isinstance(raised.value.__cause__, RuntimeError)


def test_rejected_handshake_is_not_logged_as_connected(
    caplog: pytest.LogCaptureFixture,
) -> None:
    port = find_available_port(8765)
    core = Core(CoreConfig(port=port))
    connector_uuid = uuid4()
    connector_hello = ConnectorHello(
        connector_uuid=connector_uuid,
        protocol_version=VersionDescriptor(type="alpha", major=999, minor=0, patch=0),
        connector_name="Incompatible connector",
        connector_version="0.1.0-test",
        signals=[SignalDescriptor(id=uuid4(), name="Dummy signal", type="number")],
    )

    with caplog.at_level(logging.DEBUG, logger="dynalab_core"):
        core.start()
        try:
            with socket.create_connection(("127.0.0.1", port), timeout=1) as peer:
                peer_file = peer.makefile("rb")
                assert peer_file.readline()
                peer.sendall(connector_hello.model_dump_json().encode("utf-8") + b"\n")
                assert b'"type":"handshake_rejected"' in peer_file.readline()
                rejection = _wait_for_event(caplog, "endpoint.handshake_rejected")
        finally:
            core.stop()

    assert rejection.reason == "protocol_version_mismatch"
    assert rejection.connector_uuid == str(connector_uuid)
    assert not any(
        getattr(record, "event", None) == "connector.connected"
        and getattr(record, "connector_uuid", None) == str(connector_uuid)
        for record in caplog.records
    )


def test_json_server_waits_for_connection_handlers_on_stop(
    caplog: pytest.LogCaptureFixture,
) -> None:
    port = find_available_port(8765)
    core = Core(CoreConfig(port=port))
    connector_uuid = uuid4()
    connector_hello = ConnectorHello(
        connector_uuid=connector_uuid,
        protocol_version=PROTOCOL_VERSION,
        connector_name="Test connector",
        connector_version="0.1.0-test",
        signals=[SignalDescriptor(id=uuid4(), name="Dummy signal", type="number")],
    )

    with caplog.at_level(logging.DEBUG, logger="dynalab_core"):
        core.start()
        try:
            with socket.create_connection(("127.0.0.1", port), timeout=1) as peer:
                peer_file = peer.makefile("rb")
                assert peer_file.readline()
                peer.sendall(connector_hello.model_dump_json().encode("utf-8") + b"\n")
                assert b'"type":"handshake_accepted"' in peer_file.readline()

                core.stop()
        finally:
            if core._state != "stopped":
                core.stop()

    events = [getattr(record, "event", None) for record in caplog.records]
    assert "json_server.stop_timeout" not in events
    assert "connector.connected" in events
    assert events.index("connector.registered") < events.index("connector.connected")
    assert core._json_server._stopped_event.is_set()
    assert not core._json_server._handler_tasks
    assert core._connector_registry.get(connector_uuid) is None