dynalab-core/test/test_json_server.py

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.protocols.constants import PROTOCOL_VERSION
from dynalab_core.protocols.json.errors import (
    JsonServerStartupError,
    JsonServerTimeoutError,
)
from dynalab_core.protocols.packets.handshake import ConnectorHello
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(JsonServerTimeoutError):
                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


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(JsonServerStartupError):
                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


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"not-json\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


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

    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 core._json_server._stopped_event.is_set()
    assert not core._json_server._handler_tasks
    assert core._connector_registry.get(connector_uuid) is None