Fix typing in test_wrapper

Author: dstansby

pyproject.toml

@@ -369,6 +369,7 @@ module = [
     "tests.test_store.test_logging",
     "tests.test_store.test_object",
     "tests.test_store.test_stateful",
+    "tests.test_store.test_wrapper",
 ]
 strict = false
 
@@ -377,7 +378,6 @@ strict = false
 [[tool.mypy.overrides]]
 module = [
     "tests.test_metadata.*",
-    "tests.test_store.test_wrapper",
     "tests.test_group",
     "tests.test_indexing",
     "tests.test_properties",

tests/test_store/test_wrapper.py

@@ -1,72 +1,76 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
 import pytest
 
-from zarr.core.buffer.cpu import Buffer, buffer_prototype
+from zarr.abc.store import Store
+from zarr.core.buffer import Buffer
+from zarr.core.buffer.cpu import buffer_prototype
 from zarr.storage import LocalStore, WrapperStore
 from zarr.testing.store import StoreTests
 
 if TYPE_CHECKING:
+    from pathlib import Path
+
     from _pytest.compat import LEGACY_PATH
 
-    from zarr.abc.store import Store
+    from zarr.abc.store import ByteRequest
     from zarr.core.buffer.core import BufferPrototype
 
 
 # TODO: fix this warning
 @pytest.mark.filterwarnings(
     "ignore:coroutine 'ClientCreatorContext.__aexit__' was never awaited:RuntimeWarning"
 )
-class TestWrapperStore(StoreTests[WrapperStore, Buffer]):
+class TestWrapperStore(StoreTests[WrapperStore[LocalStore], Buffer]):
     store_cls = WrapperStore
     buffer_cls = Buffer
 
-    async def get(self, store: WrapperStore, key: str) -> Buffer:
+    async def get(self, store: WrapperStore[LocalStore], key: str) -> Buffer:
         return self.buffer_cls.from_bytes((store._store.root / key).read_bytes())
 
-    async def set(self, store: WrapperStore, key: str, value: Buffer) -> None:
+    async def set(self, store: WrapperStore[LocalStore], key: str, value: Buffer) -> None:
         parent = (store._store.root / key).parent
         if not parent.exists():
             parent.mkdir(parents=True)
         (store._store.root / key).write_bytes(value.to_bytes())
 
     @pytest.fixture
-    def store_kwargs(self, tmpdir: LEGACY_PATH) -> dict[str, str]:
+    def store_kwargs(self, tmpdir: LEGACY_PATH) -> dict[str, Any]:
         return {"store": LocalStore(str(tmpdir))}
 
     @pytest.fixture
-    def open_kwargs(self, tmpdir) -> dict[str, str]:
-        return {"store_cls": LocalStore, "root": str(tmpdir)}
+    def open_kwargs(self, store_kwargs: dict[str, Any], tmp_path: Path) -> dict[str, Any]:
+        return {"store_cls": LocalStore, "root": str(tmp_path)}
 
-    def test_store_supports_writes(self, store: WrapperStore) -> None:
+    def test_store_supports_writes(self, store: WrapperStore[LocalStore]) -> None:
         assert store.supports_writes
 
-    def test_store_supports_partial_writes(self, store: WrapperStore) -> None:
+    def test_store_supports_partial_writes(self, store: WrapperStore[LocalStore]) -> None:
         assert store.supports_partial_writes
 
-    def test_store_supports_listing(self, store: WrapperStore) -> None:
+    def test_store_supports_listing(self, store: WrapperStore[LocalStore]) -> None:
         assert store.supports_listing
 
-    def test_store_repr(self, store: WrapperStore) -> None:
+    def test_store_repr(self, store: WrapperStore[LocalStore]) -> None:
         assert f"{store!r}" == f"WrapperStore(LocalStore, 'file://{store._store.root.as_posix()}')"
 
-    def test_store_str(self, store: WrapperStore) -> None:
+    def test_store_str(self, store: WrapperStore[LocalStore]) -> None:
         assert str(store) == f"wrapping-file://{store._store.root.as_posix()}"
 
-    def test_check_writeable(self, store: WrapperStore) -> None:
+    def test_check_writeable(self, store: WrapperStore[LocalStore]) -> None:
         """
         Test _check_writeable() runs without errors.
         """
         store._check_writable()
 
-    def test_close(self, store: WrapperStore) -> None:
+    def test_close(self, store: WrapperStore[LocalStore]) -> None:
         "Test store can be closed"
         store.close()
         assert not store._is_open
 
-    def test_is_open_setter_raises(self, store: WrapperStore) -> None:
+    def test_is_open_setter_raises(self, store: WrapperStore[LocalStore]) -> None:
         """
         Test that a user cannot change `_is_open` without opening the underlying store.
         """
@@ -83,7 +87,7 @@ def test_is_open_setter_raises(self, store: WrapperStore) -> None:
 @pytest.mark.parametrize("store", ["local", "memory", "zip"], indirect=True)
 async def test_wrapped_set(store: Store, capsys: pytest.CaptureFixture[str]) -> None:
     # define a class that prints when it sets
-    class NoisySetter(WrapperStore):
+    class NoisySetter(WrapperStore[Store]):
         async def set(self, key: str, value: Buffer) -> None:
             print(f"setting {key}")
             await super().set(key, value)
@@ -101,10 +105,12 @@ async def set(self, key: str, value: Buffer) -> None:
 @pytest.mark.parametrize("store", ["local", "memory", "zip"], indirect=True)
 async def test_wrapped_get(store: Store, capsys: pytest.CaptureFixture[str]) -> None:
     # define a class that prints when it sets
-    class NoisyGetter(WrapperStore):
-        def get(self, key: str, prototype: BufferPrototype) -> None:
+    class NoisyGetter(WrapperStore[Any]):
+        async def get(
+            self, key: str, prototype: BufferPrototype, byte_range: ByteRequest | None = None
+        ) -> Buffer | None:
             print(f"getting {key}")
-            return super().get(key, prototype=prototype)
+            return await super().get(key, prototype=prototype)
 
     key = "foo"
     value = Buffer.from_bytes(b"bar")