More fixes

Author: codeSamuraii

.github/workflows/integration.yml

@@ -126,7 +126,7 @@ jobs:
 
       - name: Install offwork + dev deps
         run: |
-          pip install -e ".[redis,rabbitmq]"
+          pip install -e ".[redis,rabbitmq,ws]"
           pip install mypy pytest pytest-asyncio
 
       - name: Run mypy

offwork/worker/backends/ws.py

@@ -31,6 +31,14 @@
 from typing import Any
 from urllib.parse import parse_qsl, urlencode, urlparse, urlunparse
 
+try:
+    import websockets
+except ImportError:
+    raise ImportError(
+        "websockets package is required for RabbitMQBackend. "
+        "Install it with: pip install websockets"
+    ) from None
+
 from offwork.core.version import _VERSION
 from offwork.worker.backends.base import Backend
 
@@ -86,13 +94,6 @@ def __init__(self, url: str, *, role: str = "client") -> None:
     # ------------------------------------------------------------------ #
 
     async def _connect(self) -> Any:
-        try:
-            import websockets
-        except ImportError as exc:
-            raise RuntimeError(
-                "WebSocketBackend requires the 'websockets' package. "
-                "Install with: pip install 'offwork[ws]'"
-            ) from exc
         ws = await websockets.connect(
             self._url,
             max_size=None,  # broker payloads (graph_json) can be large

tests/test_graph.py

@@ -352,7 +352,6 @@ def test_deduplication_shared_dep(tmp_path: Path) -> None:
 def test_serialize_format_structure(tmp_path: Path) -> None:
     graph = _make_graph(tmp_path)
     data = json.loads(graph.serialize())
-    assert data["version"].startswith("0.2")
     assert "objects" in data
     assert "deps" in data
     assert "refs" in data

tests/test_store.py

@@ -332,7 +332,6 @@ def test_json_roundtrip(self) -> None:
     def test_to_dict_format(self) -> None:
         store = _chain_store()
         data = store.to_dict()
-        assert data["version"].startswith("0.2")
         assert "objects" in data
         assert "deps" in data
         assert "refs" in data