ci: faster tests

Author: thearchitector

.github/workflows/ci.yaml

@@ -0,0 +1,24 @@
+name: CI
+on: [push, pull_request]
+
+jobs:
+  tests:
+    runs-on: ${{ matrix.os }}
+    strategy:
+      matrix:
+        python-version: ['3.11', '3.12', '3.13']
+        os: [ubuntu-latest, macOS-latest, windows-latest]
+
+    steps:
+      - uses: actions/checkout@v4
+      - name: Set up PDM
+        uses: pdm-project/setup-pdm@v4
+        with:
+          python-version: ${{ matrix.python-version }}
+          cache: true
+      - name: Install dependencies
+        run: |
+          pdm sync -G :all
+      - name: Run tests
+        run: |
+          pdm run -v pytest tests

.pre-commit-config.yaml

@@ -1,8 +1,8 @@
 repos:
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.11.10
+    rev: v0.12.1
     hooks:
-      - id: ruff
+      - id: ruff-check
         args:
           - --fix
           - --exit-non-zero-on-fix

datagraph/io.py

@@ -77,13 +77,13 @@ def __init__(
     async def stream(
         self, cursor: StreamCursor = StreamCursor.FROM_FIRST
     ) -> "AsyncIterator[T]":
+        client = await Supervisor.instance().client()
+
         last_entry_id: str = cursor.value
         last_stream_read: float = current_time()
 
         while True:
-            raw_message: RawStreamEntry | None = await (
-                await Supervisor.instance().client()
-            ).xread(
+            raw_message: RawStreamEntry | None = await client.xread(
                 {self._stream_key: last_entry_id},
                 options=StreamReadOptions(
                     block_ms=Supervisor.instance().config.io_read_timeout, count=1
@@ -104,6 +104,7 @@ async def stream(
                         Supervisor.instance().config.io_read_pending_timeout,
                     )
 
+                await anyio.lowlevel.checkpoint()
                 continue
 
             message = StreamEntry(raw_message)
@@ -161,15 +162,13 @@ async def write(self, value: "T") -> None:
         elif value.name != self.name:
             raise MismatchedIOError("write", self.name, value.name)
 
-        await (await Supervisor.instance().client()).xadd(
-            self._stream_key, [("ioval", self._serializer.dump(value))]
-        )
+        client = await Supervisor.instance().client()
+        await client.xadd(self._stream_key, [("ioval", self._serializer.dump(value))])
 
     async def complete(self) -> None:
-        await (await Supervisor.instance().client()).set(self._completion_key, b"true")
+        client = await Supervisor.instance().client()
+        await client.set(self._completion_key, b"true")
 
     async def is_complete(self) -> bool:
-        return (
-            await (await Supervisor.instance().client()).get(self._completion_key)
-            == b"true"
-        )
+        client = await Supervisor.instance().client()
+        return await client.get(self._completion_key) == b"true"

datagraph/serialization.py

@@ -16,22 +16,22 @@
 
 class Serializer(ABC):
     @abstractmethod
-    async def serialize(self, value: "Any") -> bytes:
+    def serialize(self, value: "Any") -> bytes:
         """Serialize a value to compressed bytestream for storage."""
         raise NotImplementedError()
 
     @abstractmethod
-    async def deserialize(self, data: bytes) -> "Any":
+    def deserialize(self, data: bytes) -> "Any":
         """Deserialize a compressed bytestream from storage."""
         raise NotImplementedError()
 
     @abstractmethod
-    async def compress(self, data: bytes) -> bytes:
+    def compress(self, data: bytes) -> bytes:
         """Compress a bytestream for storage."""
         raise NotImplementedError()
 
     @abstractmethod
-    async def decompress(self, data: bytes) -> bytes:
+    def decompress(self, data: bytes) -> bytes:
         """Decompress a bytestream from storage."""
         raise NotImplementedError()
 
@@ -58,7 +58,7 @@ def __init__(self, secret: str) -> None:
         self.secret_key: bytes = secret.encode()
 
     def serialize(self, data: "Any") -> bytes:
-        return pickletools.optimize(pickle.dumps(data, protocol=5))
+        return pickle.dumps(data, protocol=5)
 
     def deserialize(self, data: bytes) -> "Any":
         return pickle.loads(data)
@@ -78,6 +78,7 @@ def decompressor(self) -> "ZstdDecompressor":
         return self._thread_context.decompressor
 
     def compress(self, data: bytes) -> bytes:
+        data = pickletools.optimize(data)
         compressed = self.compressor.compress(data)
         signer = blake2b(digest_size=16, key=self.secret_key, usedforsecurity=True)
         signer.update(compressed)

pdm.lock

@@ -5,7 +5,7 @@
 pytest_plugins = ("celery.contrib.pytest",)
 
 
-@pytest.fixture(autouse=True)
+@pytest.fixture
 async def supervisor():
     # Use AsyncMock for the Supervisor.instance() mock so it can be awaited
     with patch(
@@ -17,7 +17,6 @@ async def supervisor():
 @pytest.fixture(
     params=[
         pytest.param(("asyncio", {"use_uvloop": False}), id="asyncio"),
-        pytest.param(("asyncio", {"use_uvloop": True}), id="asyncio+uvloop"),
         pytest.param(
             ("trio", {"restrict_keyboard_interrupt_to_checkpoints": True}), id="trio"
         ),

tests/conftest.py

@@ -180,5 +180,5 @@ async def test_execute_flow_interlaced(supervisor):
     # increasing series of timestamps, i.e. we should see CPCPCP. if C processed
     # all values before P, we'd see CCCPPP, and the sorted list would not be
     # identical to the raw zipped one
-    timestamps = list(zip(p_timestamps, c_timestamps))
+    timestamps = list(zip(p_timestamps, c_timestamps, strict=False))
     assert timestamps == sorted(timestamps)

tests/integration/__init__.py

@@ -1,6 +1,6 @@
 import pytest
 
 
-@pytest.fixture
+@pytest.fixture(scope="session")
 def celery_worker_parameters():
-    return {"shutdown_timeout": 30.0}
+    return {"shutdown_timeout": 5.0}

tests/integration/conftest.py

@@ -9,9 +9,11 @@
 from .processors import bar, consumer, first, foo, foobar, producer, second
 
 
-@pytest.fixture(params=[LocalExecutor, CeleryExecutor], ids=("local", "celery"))
-async def supervisor(request, monkeypatch):
-    monkeypatch.setattr(Supervisor, "_instance", None)
+@pytest.fixture(
+    params=[LocalExecutor, CeleryExecutor], ids=("local", "celery"), scope="module"
+)
+async def supervisor(request):
+    Supervisor._instance = None
     config = GlideClientConfiguration(addresses=[NodeAddress()], request_timeout=10000)
 
     if request.param == LocalExecutor:
@@ -35,9 +37,9 @@ async def supervisor(request, monkeypatch):
         shared_task(name=first.name, ignore_result=True)(_first)
         shared_task(name=second.name, ignore_result=True)(_second)
 
-        celery_app = request.getfixturevalue("celery_app")
+        celery_app = request.getfixturevalue("celery_session_app")
         # this is necessary to spin up the worker thread
-        _ = request.getfixturevalue("celery_worker")
+        _ = request.getfixturevalue("celery_session_worker")
 
         yield Supervisor.attach(
             glide_config=config, executor=CeleryExecutor(celery_app)

tests/integration/test_executors.py

@@ -5,17 +5,15 @@
 from datagraph.supervisor import Supervisor
 
 
-@pytest.mark.anyio
-async def test_unattached_supervisor(monkeypatch):
+def test_unattached_supervisor(monkeypatch):
     """Test that accessing Supervisor.instance() raises an error if not attached."""
     monkeypatch.setattr(Supervisor, "_instance", None)
 
     with pytest.raises(RuntimeError, match="Supervisor is not available"):
         Supervisor.instance()
 
 
-@pytest.mark.anyio
-async def test_attach_supervisor(monkeypatch):
+def test_attach_supervisor(monkeypatch):
     """Test that the attach method sets the instance correctly."""
     monkeypatch.setattr(Supervisor, "_instance", None)