@hybridlabor-api/aos 4.2.0 → 4.4.0

package/mcps/golem-rhino-mcp/tests/test_connection.py

@@ -1,543 +0,0 @@
-"""
-tests/test_connection.py
-=========================
-Unit tests for golem_3dmcp/connection.py — the RhinoConnection TCP client
-and the get_connection() singleton.
-
-All network I/O is replaced with mocks; no live socket connections are made.
-"""
-
-from __future__ import annotations
-
-import json
-import socket
-import struct
-import threading
-from unittest.mock import MagicMock, patch
-
-import pytest
-
-import golem_3dmcp.connection as _conn_module
-from golem_3dmcp.connection import (
-    RhinoCommandError,
-    RhinoConnection,
-    RhinoConnectionError,
-    RhinoTimeoutError,
-    get_connection,
-)
-
-# ---------------------------------------------------------------------------
-# Helpers
-# ---------------------------------------------------------------------------
-
-_HEADER_FORMAT = "!I"
-_HEADER_SIZE = struct.calcsize(_HEADER_FORMAT)
-
-
-def _encode(data: dict) -> bytes:
-    """Encode a dict as length-prefixed JSON (the protocol wire format)."""
-    payload = json.dumps(data, ensure_ascii=False).encode("utf-8")
-    header = struct.pack(_HEADER_FORMAT, len(payload))
-    return header + payload
-
-
-def _make_mock_socket(response: dict | None = None) -> MagicMock:
-    """
-    Create a mock socket that:
-      - ignores sendall()
-      - returns the encoded *response* from recv() calls
-    """
-    sock = MagicMock(spec=socket.socket)
-    if response is not None:
-        buf = bytearray(_encode(response))
-
-        def _recv(n: int) -> bytes:
-            chunk = bytes(buf[:n])
-            del buf[:n]
-            return chunk
-
-        sock.recv.side_effect = _recv
-    return sock
-
-
-def _success_response(result: dict, request_id: str = "test-id") -> dict:
-    return {"id": request_id, "result": result, "error": None}
-
-
-def _error_response(code: int, message: str, request_id: str = "test-id") -> dict:
-    return {
-        "id": request_id,
-        "result": None,
-        "error": {"code": code, "message": message},
-    }
-
-
-# ---------------------------------------------------------------------------
-# Fixtures
-# ---------------------------------------------------------------------------
-
-@pytest.fixture(autouse=True)
-def reset_singleton():
-    """Reset the module-level singleton before and after each test."""
-    original = _conn_module._singleton
-    _conn_module._singleton = None
-    yield
-    _conn_module._singleton = original
-
-
-@pytest.fixture
-def conn():
-    """A RhinoConnection with _sock pre-set to avoid needing connect()."""
-    c = RhinoConnection()
-    c._sock = MagicMock(spec=socket.socket)
-    return c
-
-
-# ---------------------------------------------------------------------------
-# Custom exception hierarchy
-# ---------------------------------------------------------------------------
-
-class TestExceptionHierarchy:
-
-    def test_rhino_connection_error_is_os_error(self):
-        with pytest.raises(OSError):
-            raise RhinoConnectionError("test")
-
-    def test_rhino_command_error_is_runtime_error(self):
-        with pytest.raises(RuntimeError):
-            raise RhinoCommandError(code=404, message="not found")
-
-    def test_rhino_command_error_attributes(self):
-        exc = RhinoCommandError(code=500, message="internal error")
-        assert exc.code == 500
-        assert exc.message == "internal error"
-        assert "500" in str(exc)
-        assert "internal error" in str(exc)
-
-    def test_rhino_timeout_error_is_timeout_error(self):
-        with pytest.raises(TimeoutError):
-            raise RhinoTimeoutError("timed out")
-
-
-# ---------------------------------------------------------------------------
-# connect / disconnect lifecycle
-# ---------------------------------------------------------------------------
-
-class TestConnectDisconnect:
-
-    def test_connect_creates_socket_and_pings(self):
-        """connect() should create a socket and verify it with a ping command."""
-        mock_sock = _make_mock_socket(
-            _success_response({"pong": True})
-        )
-        with patch("golem_3dmcp.connection.socket.socket", return_value=mock_sock):
-            c = RhinoConnection()
-            c.connect(host="127.0.0.1", port=9876, timeout=5)
-
-        assert c.is_connected()
-        # sendall was called (the ping request was sent)
-        assert mock_sock.sendall.called
-
-    def test_connect_raises_on_socket_error(self):
-        """If the TCP handshake fails, RhinoConnectionError is raised."""
-        mock_sock = MagicMock(spec=socket.socket)
-        mock_sock.connect.side_effect = OSError("connection refused")
-        with patch("golem_3dmcp.connection.socket.socket", return_value=mock_sock):
-            c = RhinoConnection()
-            with pytest.raises(RhinoConnectionError, match="connection refused"):
-                c.connect()
-
-    def test_connect_raises_when_ping_fails(self):
-        """If ping returns an error response, RhinoConnectionError is raised."""
-        mock_sock = _make_mock_socket(
-            _error_response(code=503, message="not ready")
-        )
-        with patch("golem_3dmcp.connection.socket.socket", return_value=mock_sock):
-            c = RhinoConnection()
-            with pytest.raises(RhinoConnectionError):
-                c.connect()
-
-    def test_disconnect_clears_socket(self, conn):
-        conn.disconnect()
-        assert not conn.is_connected()
-
-    def test_disconnect_is_idempotent(self, conn):
-        """disconnect() called multiple times should not raise."""
-        conn.disconnect()
-        conn.disconnect()  # second call — should be a no-op
-
-    def test_is_connected_false_when_no_socket(self):
-        c = RhinoConnection()
-        assert not c.is_connected()
-
-    def test_is_connected_true_when_socket_set(self, conn):
-        assert conn.is_connected()
-
-
-# ---------------------------------------------------------------------------
-# send_command — request envelope
-# ---------------------------------------------------------------------------
-
-class TestSendCommandRequestEnvelope:
-
-    def test_request_envelope_structure(self, conn):
-        """
-        send_command must build a request with 'id', 'method', 'params'
-        and send it over the wire.
-        """
-        result_payload = {"guid": "abc-123"}
-        # Set up the mock socket to return a success response.
-        buf = bytearray(_encode(_success_response(result_payload)))
-        conn._sock.recv.side_effect = lambda n: _drain(buf, n)
-
-        conn.send_command("scene.get_document_info", {"filter": None})
-
-        # Decode what was sent via sendall
-        raw = conn._sock.sendall.call_args[0][0]
-        (length,) = struct.unpack(_HEADER_FORMAT, raw[:_HEADER_SIZE])
-        request = json.loads(raw[_HEADER_SIZE:].decode("utf-8"))
-
-        assert "id" in request
-        assert len(request["id"]) > 0          # non-empty UUID
-        assert request["method"] == "scene.get_document_info"
-        assert request["params"] == {"filter": None}
-
-    def test_request_id_is_unique_per_call(self, conn):
-        """Each send_command call generates a fresh UUID4 request ID."""
-        ids = []
-
-        def capture_and_respond(raw: bytes) -> None:
-            (length,) = struct.unpack(_HEADER_FORMAT, raw[:_HEADER_SIZE])
-            req = json.loads(raw[_HEADER_SIZE:].decode("utf-8"))
-            ids.append(req["id"])
-
-        # We need two responses; queue them up.
-        responses = [
-            _encode(_success_response({})),
-            _encode(_success_response({})),
-        ]
-        combined = bytearray(b"".join(responses))
-
-        conn._sock.recv.side_effect = lambda n: _drain(combined, n)
-        conn._sock.sendall.side_effect = capture_and_respond
-
-        conn.send_command("ping", {})
-        conn.send_command("ping", {})
-
-        assert ids[0] != ids[1]
-
-    def test_send_command_returns_result(self, conn):
-        """send_command must return the 'result' dict from the response."""
-        expected = {"file_path": "/tmp/model.3dm", "units": "Meters"}
-        buf = bytearray(_encode(_success_response(expected)))
-        conn._sock.recv.side_effect = lambda n: _drain(buf, n)
-
-        result = conn.send_command("files.get_document_path", {})
-        assert result == expected
-
-    def test_send_command_returns_empty_dict_for_null_result(self, conn):
-        """If 'result' is null in the response, return an empty dict."""
-        response = {"id": "x", "result": None, "error": None}
-        buf = bytearray(_encode(response))
-        conn._sock.recv.side_effect = lambda n: _drain(buf, n)
-
-        result = conn.send_command("ping", {})
-        assert result == {}
-
-
-# ---------------------------------------------------------------------------
-# send_command — error handling
-# ---------------------------------------------------------------------------
-
-class TestSendCommandErrors:
-
-    def test_raises_rhino_command_error_on_error_response(self, conn):
-        """An error envelope from Rhino must raise RhinoCommandError."""
-        buf = bytearray(_encode(_error_response(404, "object not found")))
-        conn._sock.recv.side_effect = lambda n: _drain(buf, n)
-
-        with pytest.raises(RhinoCommandError) as exc_info:
-            conn.send_command("scene.get_object_info", {"guid": "bad"})
-
-        assert exc_info.value.code == 404
-        assert "object not found" in exc_info.value.message
-
-    def test_raises_rhino_timeout_error_on_socket_timeout(self, conn):
-        """A socket.timeout during recv must raise RhinoTimeoutError (not retried)."""
-        conn._sock.recv.side_effect = TimeoutError("timed out")
-        with pytest.raises(RhinoTimeoutError):
-            conn.send_command("slow.command", {}, timeout=1)
-
-    def test_raises_connection_error_when_not_connected(self):
-        """Calling send_command without a socket raises RhinoConnectionError."""
-        c = RhinoConnection()
-        assert c._sock is None
-        with pytest.raises(RhinoConnectionError, match="Not connected"):
-            conn_for_test = c
-            # Bypass the retry loop by calling _send_and_recv directly
-            conn_for_test._send_and_recv({"id": "x", "method": "m", "params": {}}, 5)
-
-    def test_rhino_command_error_is_not_retried(self, conn):
-        """
-        Application-level errors (RhinoCommandError) must be raised
-        immediately without triggering the auto-reconnect retry loop.
-        """
-        # Provide a single error response; if retried it would fail because
-        # there's no second response in the buffer.
-        buf = bytearray(_encode(_error_response(500, "scripting error")))
-        conn._sock.recv.side_effect = lambda n: _drain(buf, n)
-
-        with pytest.raises(RhinoCommandError):
-            conn.send_command("scripting.execute_python", {"code": "bad"})
-
-        # sendall called exactly once (no retries)
-        assert conn._sock.sendall.call_count == 1
-
-    def test_timeout_is_not_retried(self, conn):
-        """RhinoTimeoutError must propagate without retry."""
-        conn._sock.recv.side_effect = TimeoutError("timed out")
-
-        with pytest.raises(RhinoTimeoutError):
-            conn.send_command("slow.op", {})
-
-        assert conn._sock.sendall.call_count == 1
-
-    def test_malformed_response_raises_connection_error(self, conn):
-        """
-        If the response is not a dict (e.g. a JSON list), a
-        RhinoConnectionError is raised.
-        """
-        # Encode a list instead of a dict
-        payload = json.dumps([1, 2, 3]).encode("utf-8")
-        header = struct.pack(_HEADER_FORMAT, len(payload))
-        buf = bytearray(header + payload)
-        conn._sock.recv.side_effect = lambda n: _drain(buf, n)
-
-        with pytest.raises(RhinoConnectionError, match="Malformed"):
-            conn.send_command("ping", {})
-
-
-# ---------------------------------------------------------------------------
-# Auto-reconnect
-# ---------------------------------------------------------------------------
-
-class TestAutoReconnect:
-
-    def test_reconnects_on_broken_pipe_error(self):
-        """
-        On BrokenPipeError, send_command should attempt to reconnect and
-        retry the command.
-        """
-        c = RhinoConnection()
-
-        # First call: raise BrokenPipeError (simulating dropped connection).
-        # After reconnect: return a valid socket with a success response.
-        success_buf = bytearray(_encode(_success_response({"ok": True})))
-
-        reconnected_sock = MagicMock(spec=socket.socket)
-        reconnected_sock.recv.side_effect = lambda n: _drain(success_buf, n)
-
-        call_count = 0
-
-        original_send_and_recv = RhinoConnection._send_and_recv
-
-        def patched_send_and_recv(self_inner, request, timeout):
-            nonlocal call_count
-            call_count += 1
-            if call_count == 1:
-                raise BrokenPipeError("pipe broken")
-            return original_send_and_recv(self_inner, request, timeout)
-
-        def patched_reconnect(self_inner):
-            # Inject a fresh working socket
-            self_inner._sock = reconnected_sock
-
-        with patch.object(RhinoConnection, "_send_and_recv", patched_send_and_recv):
-            with patch.object(RhinoConnection, "reconnect", patched_reconnect):
-                c._sock = MagicMock(spec=socket.socket)
-                result = c.send_command("ping", {})
-
-        assert result == {"ok": True}
-        assert call_count == 2  # first attempt failed, second succeeded
-
-    def test_raises_after_max_reconnect_attempts(self):
-        """
-        If every attempt fails with BrokenPipeError, RhinoConnectionError
-        is raised after _MAX_AUTO_RECONNECT_ATTEMPTS attempts.
-        """
-        c = RhinoConnection()
-        c._sock = MagicMock(spec=socket.socket)
-
-        attempts = []
-
-        def always_broken_pipe(self_inner, request, timeout):
-            attempts.append(1)
-            raise BrokenPipeError("always broken")
-
-        def no_op_reconnect(self_inner):
-            c._sock = MagicMock(spec=socket.socket)  # give it a fresh mock
-
-        with patch.object(RhinoConnection, "_send_and_recv", always_broken_pipe):
-            with patch.object(RhinoConnection, "reconnect", no_op_reconnect):
-                with pytest.raises(RhinoConnectionError):
-                    c.send_command("ping", {})
-
-        assert len(attempts) == RhinoConnection._MAX_AUTO_RECONNECT_ATTEMPTS
-
-
-# ---------------------------------------------------------------------------
-# Thread safety
-# ---------------------------------------------------------------------------
-
-class TestThreadSafety:
-
-    def test_concurrent_send_command_calls_are_serialised(self):
-        """
-        Multiple threads calling send_command() concurrently should all
-        succeed and each receive their own response without interleaving.
-        """
-        NUM_THREADS = 8
-        results: list[dict] = []
-        errors: list[Exception] = []
-
-        def worker(thread_id: int) -> None:
-            c = RhinoConnection()
-            expected = {"thread": thread_id}
-
-            buf = bytearray(_encode(_success_response(expected)))
-            mock_sock = MagicMock(spec=socket.socket)
-            mock_sock.recv.side_effect = lambda n: _drain(buf, n)
-            c._sock = mock_sock
-
-            try:
-                result = c.send_command("ping", {})
-                results.append(result)
-            except Exception as e:
-                errors.append(e)
-
-        threads = [
-            threading.Thread(target=worker, args=(i,))
-            for i in range(NUM_THREADS)
-        ]
-        for t in threads:
-            t.start()
-        for t in threads:
-            t.join()
-
-        assert len(errors) == 0
-        assert len(results) == NUM_THREADS
-
-    def test_lock_prevents_concurrent_access_to_socket(self, conn):
-        """
-        The internal threading.Lock must block a second thread from
-        entering _send_and_recv while the first is still in progress.
-        """
-        entered_first = threading.Event()
-        can_finish = threading.Event()
-
-        success_data = _encode(_success_response({}))
-        buf = bytearray(success_data + success_data)
-        conn._sock.recv.side_effect = lambda n: _drain(buf, n)
-
-        # Patch sendall to signal thread coordination
-        def slow_sendall(data: bytes) -> None:
-            entered_first.set()
-            can_finish.wait(timeout=2)
-
-        conn._sock.sendall.side_effect = slow_sendall
-
-        thread1_result = {}
-        thread2_result = {}
-
-        def thread1():
-            thread1_result["r"] = conn.send_command("t1", {})
-
-        def thread2():
-            entered_first.wait(timeout=2)
-            can_finish.set()
-            thread2_result["r"] = conn.send_command("t2", {})
-
-        t1 = threading.Thread(target=thread1)
-        t2 = threading.Thread(target=thread2)
-        t1.start()
-        t2.start()
-        t1.join(timeout=5)
-        t2.join(timeout=5)
-
-        assert "r" in thread1_result
-        assert "r" in thread2_result
-
-
-# ---------------------------------------------------------------------------
-# get_connection() singleton
-# ---------------------------------------------------------------------------
-
-class TestGetConnectionSingleton:
-
-    def test_returns_singleton_on_second_call(self):
-        """Two calls to get_connection() with an already-connected singleton
-        should return the same instance."""
-        mock_conn = MagicMock(spec=RhinoConnection)
-        mock_conn.is_connected.return_value = True
-        _conn_module._singleton = mock_conn
-
-        result1 = get_connection()
-        result2 = get_connection()
-
-        assert result1 is result2
-
-    def test_creates_new_connection_when_singleton_is_none(self):
-        """When no singleton exists, get_connection() creates and connects one."""
-        _conn_module._singleton = None
-
-        mock_conn = MagicMock(spec=RhinoConnection)
-        mock_conn.is_connected.return_value = True
-
-        with patch("golem_3dmcp.connection.RhinoConnection", return_value=mock_conn):
-            result = get_connection(host="127.0.0.1", port=9876)
-
-        mock_conn.connect.assert_called_once_with(
-            host="127.0.0.1", port=9876, timeout=10
-        )
-        assert result is mock_conn
-
-    def test_creates_new_connection_when_existing_is_disconnected(self):
-        """
-        If the existing singleton is disconnected, get_connection() creates
-        a fresh connection.
-        """
-        stale_conn = MagicMock(spec=RhinoConnection)
-        stale_conn.is_connected.return_value = False
-        _conn_module._singleton = stale_conn
-
-        fresh_conn = MagicMock(spec=RhinoConnection)
-        fresh_conn.is_connected.return_value = True
-
-        with patch("golem_3dmcp.connection.RhinoConnection", return_value=fresh_conn):
-            result = get_connection(host="127.0.0.1", port=9876)
-
-        fresh_conn.connect.assert_called_once()
-        assert result is fresh_conn
-
-    def test_stores_new_connection_as_singleton(self):
-        """After creating a new connection it becomes the module singleton."""
-        _conn_module._singleton = None
-
-        mock_conn = MagicMock(spec=RhinoConnection)
-        mock_conn.is_connected.return_value = True
-
-        with patch("golem_3dmcp.connection.RhinoConnection", return_value=mock_conn):
-            get_connection(host="127.0.0.1", port=9876)
-
-        assert _conn_module._singleton is mock_conn
-
-
-# ---------------------------------------------------------------------------
-# Helpers used within this module
-# ---------------------------------------------------------------------------
-
-def _drain(buf: bytearray, n: int) -> bytes:
-    """Pop and return the first *n* bytes from *buf*."""
-    chunk = bytes(buf[:n])
-    del buf[:n]
-    return chunk