test: add tests for SupportsSetRange on MemoryStore and LocalStore

Tests cover isinstance check, async set_range, sync set_range_sync,
and edge case (writing at end of value).

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: d-v-b

src/zarr/abc/store.py

@@ -712,7 +712,15 @@ async def set_if_not_exists(self, default: Buffer) -> None: ...
 
 @runtime_checkable
 class SupportsSetRange(Protocol):
-    """Protocol for stores that support writing to a byte range within an existing value."""
+    """Protocol for stores that support writing to a byte range within an existing value.
+
+    Overwrites ``len(value)`` bytes starting at byte offset ``start`` within the
+    existing stored value for ``key``. The key must already exist and the write
+    must fit within the existing value (i.e., ``start + len(value) <= len(existing)``).
+
+    Behavior when the write extends past the end of the existing value is
+    implementation-specific and should not be relied upon.
+    """
 
     async def set_range(self, key: str, value: Buffer, start: int) -> None: ...
 

tests/test_store/test_local.py

@@ -10,6 +10,7 @@
 
 import zarr
 from zarr import create_array
+from zarr.abc.store import SupportsSetRange
 from zarr.core.buffer import Buffer, cpu
 from zarr.core.sync import sync
 from zarr.storage import LocalStore
@@ -162,6 +163,54 @@ def test_get_json_sync_with_prototype_none(
         result = store._get_json_sync(key, prototype=buffer_cls)
         assert result == data
 
+    def test_supports_set_range(self, store: LocalStore) -> None:
+        """LocalStore should implement SupportsSetRange."""
+        assert isinstance(store, SupportsSetRange)
+
+    @pytest.mark.parametrize(
+        ("start", "patch", "expected"),
+        [
+            (0, b"XX", b"XXAAAAAAAA"),
+            (3, b"XX", b"AAAXXAAAAA"),
+            (8, b"XX", b"AAAAAAAAXX"),
+            (0, b"ZZZZZZZZZZ", b"ZZZZZZZZZZ"),
+            (5, b"B", b"AAAAABAAAA"),
+            (0, b"BCDE", b"BCDEAAAAAA"),
+        ],
+        ids=["start", "middle", "end", "full-overwrite", "single-byte", "multi-byte-start"],
+    )
+    async def test_set_range(
+        self, store: LocalStore, start: int, patch: bytes, expected: bytes
+    ) -> None:
+        """set_range should overwrite bytes at the given offset."""
+        await store.set("test/key", cpu.Buffer.from_bytes(b"AAAAAAAAAA"))
+        await store.set_range("test/key", cpu.Buffer.from_bytes(patch), start=start)
+        result = await store.get("test/key", prototype=cpu.buffer_prototype)
+        assert result is not None
+        assert result.to_bytes() == expected
+
+    @pytest.mark.parametrize(
+        ("start", "patch", "expected"),
+        [
+            (0, b"XX", b"XXAAAAAAAA"),
+            (3, b"XX", b"AAAXXAAAAA"),
+            (8, b"XX", b"AAAAAAAAXX"),
+            (0, b"ZZZZZZZZZZ", b"ZZZZZZZZZZ"),
+            (5, b"B", b"AAAAABAAAA"),
+            (0, b"BCDE", b"BCDEAAAAAA"),
+        ],
+        ids=["start", "middle", "end", "full-overwrite", "single-byte", "multi-byte-start"],
+    )
+    def test_set_range_sync(
+        self, store: LocalStore, start: int, patch: bytes, expected: bytes
+    ) -> None:
+        """set_range_sync should overwrite bytes at the given offset."""
+        sync(store.set("test/key", cpu.Buffer.from_bytes(b"AAAAAAAAAA")))
+        store.set_range_sync("test/key", cpu.Buffer.from_bytes(patch), start=start)
+        result = store.get_sync(key="test/key", prototype=cpu.buffer_prototype)
+        assert result is not None
+        assert result.to_bytes() == expected
+
 
 @pytest.mark.parametrize("exclusive", [True, False])
 def test_atomic_write_successful(tmp_path: pathlib.Path, exclusive: bool) -> None:

tests/test_store/test_memory.py

@@ -9,6 +9,7 @@
 import pytest
 
 import zarr
+from zarr.abc.store import SupportsSetRange
 from zarr.core.buffer import Buffer, cpu, gpu
 from zarr.core.sync import sync
 from zarr.errors import ZarrUserWarning
@@ -127,6 +128,55 @@ def test_get_json_sync_with_prototype_none(
         result = store._get_json_sync(key, prototype=buffer_cls)
         assert result == data
 
+    def test_supports_set_range(self, store: MemoryStore) -> None:
+        """MemoryStore should implement SupportsSetRange."""
+        assert isinstance(store, SupportsSetRange)
+
+    @pytest.mark.parametrize(
+        ("start", "patch", "expected"),
+        [
+            (0, b"XX", b"XXAAAAAAAA"),
+            (3, b"XX", b"AAAXXAAAAA"),
+            (8, b"XX", b"AAAAAAAAXX"),
+            (0, b"ZZZZZZZZZZ", b"ZZZZZZZZZZ"),
+            (5, b"B", b"AAAAABAAAA"),
+            (0, b"BCDE", b"BCDEAAAAAA"),
+        ],
+        ids=["start", "middle", "end", "full-overwrite", "single-byte", "multi-byte-start"],
+    )
+    async def test_set_range(
+        self, store: MemoryStore, start: int, patch: bytes, expected: bytes
+    ) -> None:
+        """set_range should overwrite bytes at the given offset."""
+        await store.set("test/key", cpu.Buffer.from_bytes(b"AAAAAAAAAA"))
+        await store.set_range("test/key", cpu.Buffer.from_bytes(patch), start=start)
+        result = await store.get("test/key", prototype=cpu.buffer_prototype)
+        assert result is not None
+        assert result.to_bytes() == expected
+
+    @pytest.mark.parametrize(
+        ("start", "patch", "expected"),
+        [
+            (0, b"XX", b"XXAAAAAAAA"),
+            (3, b"XX", b"AAAXXAAAAA"),
+            (8, b"XX", b"AAAAAAAAXX"),
+            (0, b"ZZZZZZZZZZ", b"ZZZZZZZZZZ"),
+            (5, b"B", b"AAAAABAAAA"),
+            (0, b"BCDE", b"BCDEAAAAAA"),
+        ],
+        ids=["start", "middle", "end", "full-overwrite", "single-byte", "multi-byte-start"],
+    )
+    def test_set_range_sync(
+        self, store: MemoryStore, start: int, patch: bytes, expected: bytes
+    ) -> None:
+        """set_range_sync should overwrite bytes at the given offset."""
+        store._is_open = True
+        store._store_dict["test/key"] = cpu.Buffer.from_bytes(b"AAAAAAAAAA")
+        store.set_range_sync("test/key", cpu.Buffer.from_bytes(patch), start=start)
+        result = store.get_sync(key="test/key", prototype=cpu.buffer_prototype)
+        assert result is not None
+        assert result.to_bytes() == expected
+
 
 # TODO: fix this warning
 @pytest.mark.filterwarnings("ignore:Unclosed client session:ResourceWarning")