Stop re-exporting websocket classes from s2python.connection

examples/async_frbc_rm.py

@@ -34,7 +34,8 @@
     FRBCActuatorStatus,
 )
 from s2python.connection import AssetDetails
-from s2python.connection.async_ import S2AsyncConnection, WebsocketClientMedium
+from s2python.connection.async_ import S2AsyncConnection
+from s2python.connection.async_.medium.websocket import WebsocketClientMedium
 from s2python.connection.async_.control_type.class_based import (
     FRBCControlType,
     NoControlControlType,

examples/quickstart_websocket_rm.py

@@ -29,7 +29,8 @@
     FRBCStorageStatus,
     FRBCActuatorStatus,
 )
-from s2python.connection import AssetDetails, BlockingWebsocketClientRM
+from s2python.connection import AssetDetails
+from s2python.connection.quickstarts import BlockingWebsocketClientRM
 from s2python.connection.sync import S2SyncConnection
 from s2python.connection.sync.control_type.class_based import FRBCControlType, NoControlControlType
 

examples/sync_frbc_rm.py

@@ -36,7 +36,7 @@
 )
 from s2python.connection import AssetDetails
 from s2python.connection.sync import S2SyncConnection
-from s2python.connection.async_ import WebsocketClientMedium
+from s2python.connection.async_.medium.websocket import WebsocketClientMedium
 from s2python.connection.sync.control_type.class_based import (
     FRBCControlType,
     NoControlControlType,

src/s2python/connection/__init__.py

@@ -2,7 +2,6 @@
 from s2python.connection.connection_events import ConnectionStarted, ConnectionStopped, S2ConnectionEvent
 from s2python.connection.control_type import RoleHandler
 from s2python.connection.types import S2ConnectionEventsAndMessages
-from s2python.connection.quickstarts import BlockingWebsocketClientRM
 
 __all__ = [
     "AssetDetails",
@@ -11,5 +10,4 @@
     "S2ConnectionEvent",
     "RoleHandler",
     "S2ConnectionEventsAndMessages",
-    "BlockingWebsocketClientRM"
 ]

src/s2python/connection/async_/__init__.py

@@ -1,13 +1,11 @@
 from s2python.connection.async_.connection import S2AsyncConnection, CouldNotReceiveStatusReceptionError
 from s2python.connection.async_.message_handlers import S2EventHandlerAsync
 from s2python.connection.async_.medium.s2_medium import S2MediumConnection
-from s2python.connection.async_.medium.websocket import WebsocketClientMedium
 
 
 __all__ = [
     "S2AsyncConnection",
     "CouldNotReceiveStatusReceptionError",
     "S2EventHandlerAsync",
     "S2MediumConnection",
-    "WebsocketClientMedium"
 ]

src/s2python/connection/quickstarts.py

@@ -3,7 +3,7 @@
 import threading
 
 from s2python.connection.asset_details import AssetDetails
-from s2python.connection.async_ import WebsocketClientMedium
+from s2python.connection.async_.medium.websocket import WebsocketClientMedium
 from s2python.connection.sync import S2SyncConnection
 from s2python.connection.sync.control_type.class_based import ResourceManagerHandler, S2ControlType
 

tests/unit/connection/init_test.py

@@ -0,0 +1,62 @@
+"""Tests that s2python.connection does not require the optional 'websockets' dependency."""
+
+import os
+import subprocess
+import sys
+
+import pytest
+
+SRC_DIR = os.path.join(os.path.dirname(__file__), "..", "..", "..", "src")
+
+
+def _run_without_websockets(code: str) -> subprocess.CompletedProcess[str]:
+    # Setting sys.modules['websockets'] to None makes any import of it raise ImportError.
+    script = "import sys\nsys.modules['websockets'] = None\n" + code
+    python_path = [SRC_DIR]
+    if os.environ.get("PYTHONPATH"):
+        python_path.append(os.environ["PYTHONPATH"])
+    env = dict(os.environ, PYTHONPATH=os.pathsep.join(python_path))
+    return subprocess.run(
+        [sys.executable, "-c", script], capture_output=True, text=True, env=env, check=False
+    )
+
+
+@pytest.mark.parametrize(
+    "import_statement",
+    [
+        "import s2python.connection",
+        "from s2python.connection import AssetDetails",
+        "import s2python.connection.async_",
+        "import s2python.connection.sync",
+        "from s2python.connection.async_.control_type.class_based import OMBCControlType",
+        "from s2python.connection.sync.control_type.class_based import OMBCControlType",
+    ],
+)
+def test__import__works_without_websockets(import_statement: str) -> None:
+    # Arrange
+    code = import_statement + "\nassert 's2python.connection.async_.medium.websocket' not in sys.modules\n"
+
+    # Act
+    result = _run_without_websockets(code)
+
+    # Assert
+    assert result.returncode == 0, result.stderr
+
+
+@pytest.mark.parametrize(
+    "import_statement",
+    [
+        "from s2python.connection.quickstarts import BlockingWebsocketClientRM",
+        "from s2python.connection.async_.medium.websocket import WebsocketClientMedium",
+    ],
+)
+def test__import__websocket_modules_require_websockets(import_statement: str) -> None:
+    # Arrange
+    code = import_statement
+
+    # Act
+    result = _run_without_websockets(code)
+
+    # Assert
+    assert result.returncode != 0
+    assert "The 'websockets' package is required" in result.stderr