feat: 基于Python ROS2的控制程序

python/tests/test_sessions.py

@@ -0,0 +1,180 @@
+from __future__ import annotations
+
+from contextlib import contextmanager
+from pathlib import Path
+import socket
+import subprocess
+import sys
+import threading
+import time
+
+import pytest
+
+
+pytestmark = pytest.mark.skipif(sys.platform != 'linux', reason='Linux-only OmniSocket extension')
+
+ROOT = Path(__file__).resolve().parents[2]
+PYTHON_ROOT = ROOT / 'python'
+if str(PYTHON_ROOT) not in sys.path:
+    sys.path.insert(0, str(PYTHON_ROOT))
+
+omnisocket = pytest.importorskip('omnisocket')
+
+CONTROL_DEFAULTS = omnisocket.CONTROL_DEFAULTS
+MSG_TYPE_BINARY = omnisocket.MSG_TYPE_BINARY
+Session = omnisocket.Session
+UdpSession = omnisocket.UdpSession
+
+
+def _reserve_port() -> int:
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
+        sock.bind(('127.0.0.1', 0))
+        return int(sock.getsockname()[1])
+
+
+@contextmanager
+def _run_server(binary_name: str, listen_addr: str):
+    binary = ROOT / 'bin' / binary_name
+    if not binary.exists():
+        pytest.skip(f'{binary} is not built')
+
+    process = subprocess.Popen(
+        [str(binary), '-listen', listen_addr],
+        cwd=str(ROOT),
+        stdout=subprocess.DEVNULL,
+        stderr=subprocess.DEVNULL,
+    )
+    try:
+        time.sleep(0.2)
+        yield process
+    finally:
+        process.terminate()
+        try:
+            process.wait(timeout=2.0)
+        except subprocess.TimeoutExpired:
+            process.kill()
+            process.wait(timeout=2.0)
+
+
+def _connect_with_retry(session_cls, *, transport: str, server_addr: str, peer_id: str):
+    deadline = time.monotonic() + 3.0
+    last_error: Exception | None = None
+
+    while time.monotonic() < deadline:
+        session = session_cls()
+        try:
+            kwargs: dict[str, object] = {
+                'server_addr': server_addr,
+                'peer_id': peer_id,
+            }
+            if transport == 'kcp':
+                kwargs.update(CONTROL_DEFAULTS)
+            else:
+                kwargs['enable_timestamping'] = False
+            session.connect(**kwargs)
+            return session
+        except OSError as exc:
+            last_error = exc
+            time.sleep(0.1)
+
+    raise AssertionError(f'failed to connect {peer_id} to {server_addr}: {last_error}')
+
+
+@pytest.mark.parametrize(
+    ('transport', 'binary_name', 'session_cls'),
+    [
+        ('udp', 'udpserver', UdpSession),
+        ('kcp', 'kcpserver', Session),
+    ],
+)
+def test_control_sessions_smoke(transport: str, binary_name: str, session_cls) -> None:
+    port = _reserve_port()
+    listen_addr = f'127.0.0.1:{port}'
+    sender_id = f'pytest-{transport}-sender'
+    receiver_id = f'pytest-{transport}-receiver'
+
+    with _run_server(binary_name, listen_addr):
+        sender = _connect_with_retry(session_cls, transport=transport, server_addr=listen_addr, peer_id=sender_id)
+        receiver = _connect_with_retry(session_cls, transport=transport, server_addr=listen_addr, peer_id=receiver_id)
+
+        try:
+            assert receiver.recv(timeout_ms=20) is None
+
+            payload = b'control-packet-1'
+            sender.send(to=receiver_id, data=payload)
+            from_peer, msg_type, recv_payload = receiver.recv(timeout_ms=1000)
+            assert from_peer == sender_id
+            assert msg_type == MSG_TYPE_BINARY
+            assert recv_payload == payload
+
+            payload2 = b'control-packet-2'
+            sender.send(to=receiver_id, data=payload2)
+            recv_buffer = bytearray(128)
+            meta = receiver.recv_into(buffer=recv_buffer, timeout_ms=1000)
+            assert meta is not None
+            assert meta['from'] == sender_id
+            assert meta['msg_type'] == MSG_TYPE_BINARY
+            assert meta['body_len'] == len(payload2)
+            assert bytes(recv_buffer[: meta['body_len']]) == payload2
+
+            sender_stats = sender.stats()
+            receiver_stats = receiver.stats()
+            assert sender_stats['connected'] == 1
+            assert receiver_stats['connected'] == 1
+            assert sender_stats['send_calls'] >= 2
+            assert receiver_stats['recv_calls'] >= 2
+        finally:
+            sender.close()
+            receiver.close()
+
+
+def test_udp_session_close_interrupts_blocking_recv() -> None:
+    port = _reserve_port()
+    listen_addr = f'127.0.0.1:{port}'
+    receiver_id = 'pytest-udp-blocking-recv'
+
+    with _run_server('udpserver', listen_addr):
+        receiver = _connect_with_retry(
+            UdpSession,
+            transport='udp',
+            server_addr=listen_addr,
+            peer_id=receiver_id,
+        )
+
+        recv_error: list[BaseException] = []
+        close_error: list[BaseException] = []
+        recv_started = threading.Event()
+        recv_done = threading.Event()
+        close_done = threading.Event()
+
+        def recv_worker() -> None:
+            recv_started.set()
+            try:
+                receiver.recv()
+            except BaseException as exc:  # pragma: no cover - assertion is on thread completion
+                recv_error.append(exc)
+            finally:
+                recv_done.set()
+
+        def close_worker() -> None:
+            try:
+                receiver.close()
+            except BaseException as exc:  # pragma: no cover - assertion is on thread completion
+                close_error.append(exc)
+            finally:
+                close_done.set()
+
+        recv_thread = threading.Thread(target=recv_worker, daemon=True)
+        recv_thread.start()
+        assert recv_started.wait(timeout=1.0)
+        time.sleep(0.05)
+
+        close_thread = threading.Thread(target=close_worker, daemon=True)
+        close_thread.start()
+
+        assert close_done.wait(timeout=1.0), 'UdpSession.close() blocked while recv() was waiting'
+        assert recv_done.wait(timeout=1.0), 'UdpSession.recv() stayed blocked after close()'
+        assert not close_thread.is_alive()
+        assert not recv_thread.is_alive()
+        assert not close_error
+        assert not recv_error or isinstance(recv_error[0], OSError)