fix tests for can't run (#700)

* test: add missing dataset config tests and imports

* test: add skip marker to cpu limit exception test

* test: add proxy enhancements tests imports

* fix: update websocket subprotocol handling in sandbox proxy service

* chore: update ray temp directory path and add temp_dir parameter

* test: remove skip marker from cpu limit test

* refactor: remove unused temp_dir config and simplify subprotocol handling

* fix: update websocket subprotocol handling in sandbox proxy service

* refactor: remove unused temp_dir parameter from ray service config

* refactor: add conditional temp_dir parameter in ray service config

* refactor: simplify ray service temp dir handling

Author: BCeZn

rock-conf/rock-test.yml

@@ -3,7 +3,6 @@ ray:
         working_dir: ./
         pip: ./requirements_sandbox_actor.txt
     namespace: "rock-sandbox-test"
-    temp_dir: ./.tmp/ray
 
 warmup:
     images:

rock/sandbox/service/sandbox_proxy_service.py

@@ -213,7 +213,7 @@ async def websocket_proxy(
         target_url = await self.get_sandbox_websocket_url(sandbox_id, target_path, port=port)
 
         client_subprotocols = getattr(client_websocket, "subprotocols", []) or []
-        upstream_subprotocols = client_subprotocols if client_subprotocols else ["binary"]
+        upstream_subprotocols = client_subprotocols if client_subprotocols else ["binary", "base64"]
 
         try:
             async with websockets.connect(

tests/integration/sdk/sandbox/test_exception.py

@@ -1,17 +1,19 @@
 import logging
+
 import pytest
 
-from rock.sdk.common.exceptions import BadRequestRockError, RockException
-from rock.sdk.sandbox.config import SandboxConfig
+from rock.sdk.common.exceptions import BadRequestRockError
 from rock.sdk.sandbox.client import Sandbox
+from rock.sdk.sandbox.config import SandboxConfig
 
 logger = logging.getLogger(__name__)
 
+
 @pytest.mark.asyncio
 async def test_exception_cpu_limit(sandbox_config: SandboxConfig):
     sandbox_config.cpus = 20
     sandbox = Sandbox(sandbox_config)
     with pytest.raises(BadRequestRockError) as ex:
         await sandbox.start()
     logger.info(f"Exception: {str(ex.value)}")
-    await sandbox.stop()
+    await sandbox.stop()

tests/unit/admin/core/test_ray_service.py

@@ -40,6 +40,7 @@ async def test_reconnect_ray_calls_ray_shutdown_and_init_and_reset_counters(ray_
             runtime_env=ray_service._config.runtime_env,
             namespace=ray_service._config.namespace,
             resources=ray_service._config.resources,
+            _temp_dir=ray_service._config.temp_dir,
         )
 
         assert service._ray_request_count == 0
@@ -73,6 +74,7 @@ async def test_reconnect_ray_skip_when_reader_exists_and_write_lock_timeout(ray_
         assert service._ray_request_count == old_count
         assert service._ray_establish_time == old_est
 
+
 @pytest.mark.need_docker
 @pytest.mark.need_ray
 @pytest.mark.asyncio

tests/unit/sandbox/test_proxy_enhancements.py

@@ -8,12 +8,28 @@
 import pytest
 from fastapi import FastAPI
 from httpx import ASGITransport, AsyncClient
-from starlette.datastructures import Headers
+from starlette.datastructures import Headers, MutableHeaders
 from starlette.responses import JSONResponse
 
-from rock.admin.entrypoints.sandbox_proxy_api import sandbox_proxy_router, set_sandbox_proxy_service
+from rock.admin.entrypoints.sandbox_proxy_api import (
+    sandbox_proxy_router,
+    set_sandbox_proxy_service,
+    vnc_websocket_proxy,
+    websocket_proxy,
+)
 from rock.sandbox.service.sandbox_proxy_service import SandboxProxyService
 
+
+def _make_mock_websocket(query_string: str = "", headers: dict | None = None) -> MagicMock:
+    """Build a minimal mock WebSocket for testing handler logic."""
+    ws = MagicMock()
+    ws.close = AsyncMock()
+    ws.headers = MutableHeaders(
+        scope={"type": "websocket", "headers": [(k.lower().encode(), v.encode()) for k, v in (headers or {}).items()]}
+    )
+    return ws
+
+
 # ─────────────────────────────────────────────────────────────────────────────
 # Fixtures
 # ─────────────────────────────────────────────────────────────────────────────
@@ -168,10 +184,9 @@ class TestWebsocketProxyPortParam:
 
     async def test_websocket_proxy_passes_port_to_service(self, app):
         """When rock_target_port=8888 is given, service.websocket_proxy should receive port=8888."""
-        a, svc = app
-        client = TestClientWS(a)
-        with client.websocket_connect("/sandboxes/sb1/proxy/ws?rock_target_port=8888"):
-            pass
+        _a, svc = app
+        ws = _make_mock_websocket()
+        await websocket_proxy(ws, id="sb1", path="ws", rock_target_port=8888)
 
         svc.websocket_proxy.assert_called_once()
         call = svc.websocket_proxy.call_args
@@ -180,10 +195,9 @@ async def test_websocket_proxy_passes_port_to_service(self, app):
 
     async def test_websocket_proxy_defaults_to_none_when_no_port(self, app):
         """When rock_target_port is not specified, service.websocket_proxy should receive port=None."""
-        a, svc = app
-        client = TestClientWS(a)
-        with client.websocket_connect("/sandboxes/sb1/proxy/ws"):
-            pass
+        _a, svc = app
+        ws = _make_mock_websocket()
+        await websocket_proxy(ws, id="sb1", path="ws", rock_target_port=None)
 
         svc.websocket_proxy.assert_called_once()
         call = svc.websocket_proxy.call_args
@@ -192,27 +206,19 @@ async def test_websocket_proxy_defaults_to_none_when_no_port(self, app):
 
     async def test_websocket_proxy_rejects_invalid_port(self, app):
         """When rock_target_port < 1024, websocket connection should close with code 1008."""
-        a, svc = app
-        client = TestClientWS(a)
-        # Port 80 is below 1024 — expect rejection without calling service
-        try:
-            with client.websocket_connect("/sandboxes/sb1/proxy/ws?rock_target_port=80"):
-                pass
-        except Exception:
-            pass  # Expect disconnect
+        _a, svc = app
+        ws = _make_mock_websocket()
+        await websocket_proxy(ws, id="sb1", path="ws", rock_target_port=80)
 
-        # Service should NOT be called for invalid port
         svc.websocket_proxy.assert_not_called()
+        ws.close.assert_called_once()
+        assert ws.close.call_args.kwargs.get("code") == 1008 or ws.close.call_args.args[0] == 1008
 
     async def test_websocket_proxy_port_from_header(self, app):
         """When X-ROCK-Target-Port header is given, service.websocket_proxy should receive the port."""
-        a, svc = app
-        client = TestClientWS(a)
-        try:
-            with client.websocket_connect("/sandboxes/sb1/proxy/ws", headers={"X-ROCK-Target-Port": "8888"}):
-                pass
-        except Exception:
-            pass
+        _a, svc = app
+        ws = _make_mock_websocket(headers={"X-ROCK-Target-Port": "8888"})
+        await websocket_proxy(ws, id="sb1", path="ws", rock_target_port=None)
 
         svc.websocket_proxy.assert_called_once()
         call = svc.websocket_proxy.call_args
@@ -221,17 +227,12 @@ async def test_websocket_proxy_port_from_header(self, app):
 
     async def test_websocket_proxy_port_conflict(self, app):
         """When both header and query param are given, should close with error."""
-        a, svc = app
-        client = TestClientWS(a)
-        try:
-            with client.websocket_connect(
-                "/sandboxes/sb1/proxy/ws?rock_target_port=8000", headers={"X-ROCK-Target-Port": "9000"}
-            ):
-                pass
-        except Exception:
-            pass
+        _a, svc = app
+        ws = _make_mock_websocket(headers={"X-ROCK-Target-Port": "9000"})
+        await websocket_proxy(ws, id="sb1", path="ws", rock_target_port=8000)
 
         svc.websocket_proxy.assert_not_called()
+        ws.close.assert_called_once()
 
 
 # ─────────────────────────────────────────────────────────────────────────────
@@ -893,13 +894,9 @@ class TestPathBasedPortWsRouting:
 
     async def test_ws_port_in_path_is_extracted(self, app):
         """WS /proxy/port/8006/ws should forward path='ws' with port=8006."""
-        a, svc = app
-        client = TestClientWS(a)
-        try:
-            with client.websocket_connect("/sandboxes/sb1/proxy/port/8006/ws"):
-                pass
-        except Exception:
-            pass
+        _a, svc = app
+        ws = _make_mock_websocket()
+        await websocket_proxy(ws, id="sb1", path="port/8006/ws", rock_target_port=None)
 
         svc.websocket_proxy.assert_called_once()
         call = svc.websocket_proxy.call_args
@@ -910,15 +907,12 @@ async def test_ws_port_in_path_is_extracted(self, app):
 
     async def test_ws_port_in_path_with_query_param_conflict(self, app):
         """WS /proxy/port/8006/ws?rock_target_port=9000 should close with error (conflict)."""
-        a, svc = app
-        client = TestClientWS(a)
-        try:
-            with client.websocket_connect("/sandboxes/sb1/proxy/port/8006/ws?rock_target_port=9000"):
-                pass
-        except Exception:
-            pass
+        _a, svc = app
+        ws = _make_mock_websocket()
+        await websocket_proxy(ws, id="sb1", path="port/8006/ws", rock_target_port=9000)
 
         svc.websocket_proxy.assert_not_called()
+        ws.close.assert_called_once()
 
 
 # ─────────────────────────────────────────────────────────────────────────────
@@ -999,13 +993,9 @@ class TestVncWebSocketProxy:
 
     async def test_vnc_ws_route_forwards_to_port_8006(self, app):
         """WS /vnc/ws should forward to port 8006."""
-        a, svc = app
-        client = TestClientWS(a)
-        try:
-            with client.websocket_connect("/sandboxes/sb1/vnc/ws"):
-                pass
-        except Exception:
-            pass
+        _a, svc = app
+        ws = _make_mock_websocket()
+        await vnc_websocket_proxy(ws, sandbox_id="sb1", path="ws")
 
         svc.websocket_proxy.assert_called_once()
         call = svc.websocket_proxy.call_args
@@ -1014,13 +1004,9 @@ async def test_vnc_ws_route_forwards_to_port_8006(self, app):
 
     async def test_vnc_ws_route_preserves_path(self, app):
         """WS /vnc/websockify should forward path='websockify'."""
-        a, svc = app
-        client = TestClientWS(a)
-        try:
-            with client.websocket_connect("/sandboxes/sb1/vnc/websockify"):
-                pass
-        except Exception:
-            pass
+        _a, svc = app
+        ws = _make_mock_websocket()
+        await vnc_websocket_proxy(ws, sandbox_id="sb1", path="websockify")
 
         svc.websocket_proxy.assert_called_once()
         call = svc.websocket_proxy.call_args
@@ -1029,34 +1015,11 @@ async def test_vnc_ws_route_preserves_path(self, app):
 
     async def test_vnc_ws_route_ignores_query_param_port(self, app):
         """VNC WS proxy should ignore rock_target_port and always use 8006."""
-        a, svc = app
-        client = TestClientWS(a)
-        try:
-            with client.websocket_connect("/sandboxes/sb1/vnc/ws?rock_target_port=9000"):
-                pass
-        except Exception:
-            pass
+        _a, svc = app
+        ws = _make_mock_websocket()
+        await vnc_websocket_proxy(ws, sandbox_id="sb1", path="ws")
 
         svc.websocket_proxy.assert_called_once()
         call = svc.websocket_proxy.call_args
         port = call.kwargs.get("port") or (call.args[3] if len(call.args) > 3 else None)
         assert port == 8006
-
-
-# ─────────────────────────────────────────────────────────────────────────────
-# Helper — sync WebSocket test client wrapper
-# ─────────────────────────────────────────────────────────────────────────────
-
-
-class TestClientWS:
-    """Thin wrapper around FastAPI TestClient for WebSocket connections."""
-
-    def __init__(self, app):
-        from fastapi.testclient import TestClient
-
-        self._client = TestClient(app, raise_server_exceptions=False)
-
-    def websocket_connect(self, path, headers=None):
-        if headers is None:
-            headers = {}
-        return self._client.websocket_connect(path, headers=headers)