Add hypothesis testing

Author: jsignell

properties/test_parallelcompat.py

@@ -0,0 +1,71 @@
+import numpy as np
+import pytest
+
+pytest.importorskip("hypothesis")
+# isort: split
+
+from hypothesis import given
+
+import xarray.testing.strategies as xrst
+from xarray.namedarray.parallelcompat import ChunkManagerEntrypoint
+
+
+class TestPreserveChunks:
+    @given(xrst.shape_and_chunks())
+    def test_preserve_all_chunks(
+        self, shape_and_chunks: tuple[tuple[int, ...], tuple[int, ...]]
+    ) -> None:
+        shape, previous_chunks = shape_and_chunks
+        typesize = 8
+        target = 1024 * 1024
+
+        actual = ChunkManagerEntrypoint.preserve_chunks(
+            chunks=("preserve",) * len(shape),
+            shape=shape,
+            target=target,
+            typesize=typesize,
+            previous_chunks=previous_chunks,
+        )
+        for i, chunk in enumerate(actual):
+            if chunk != shape[i]:
+                assert chunk >= previous_chunks[i]
+                assert chunk % previous_chunks[i] == 0
+                assert chunk <= shape[i]
+
+        if actual != shape:
+            assert np.prod(actual) * typesize >= 0.5 * target
+
+    @pytest.mark.parametrize("first_chunk", [-1, (), 1])
+    @given(xrst.shape_and_chunks(min_dims=2))
+    def test_preserve_some_chunks(
+        self,
+        first_chunk: int | tuple[int, ...],
+        shape_and_chunks: tuple[tuple[int, ...], tuple[int, ...]],
+    ) -> None:
+        shape, previous_chunks = shape_and_chunks
+        typesize = 4
+        target = 2 * 1024 * 1024
+
+        actual = ChunkManagerEntrypoint.preserve_chunks(
+            chunks=(first_chunk, *["preserve" for _ in range(len(shape) - 1)]),
+            shape=shape,
+            target=target,
+            typesize=typesize,
+            previous_chunks=previous_chunks,
+        )
+        for i, chunk in enumerate(actual):
+            if i == 0:
+                if first_chunk == 1:
+                    assert chunk == 1
+                elif first_chunk == -1:
+                    assert chunk == shape[i]
+                elif first_chunk == ():
+                    assert chunk == previous_chunks[i]
+            elif chunk != shape[i]:
+                assert chunk >= previous_chunks[i]
+                assert chunk % previous_chunks[i] == 0
+                assert chunk <= shape[i]
+
+        # if we have more than one chunk, make sure the chunks are big enough
+        if actual[1:] != shape[1:]:
+            assert np.prod(actual) * typesize >= 0.5 * target

xarray/testing/strategies.py

@@ -31,6 +31,7 @@
     "names",
     "outer_array_indexers",
     "pandas_index_dtypes",
+    "shape_and_chunks",
     "supported_dtypes",
     "unique_subset_of",
     "variables",
@@ -210,6 +211,73 @@ def dimension_sizes(
     )
 
 
+@st.composite
+def shape_and_chunks(
+    draw: st.DrawFn,
+    *,
+    min_dims: int = 1,
+    max_dims: int = 4,
+    min_size: int = 1,
+    max_size: int = 900,
+) -> tuple[tuple[int, ...], tuple[int, ...]]:
+    """
+    Generate a shape tuple and corresponding chunks tuple.
+
+    Each element in the chunks tuple is smaller than or equal to the
+    corresponding element in the shape tuple.
+
+    Requires the hypothesis package to be installed.
+
+    Parameters
+    ----------
+    min_dims : int, optional
+        Minimum number of dimensions. Default is 1.
+    max_dims : int, optional
+        Maximum number of dimensions. Default is 4.
+    min_size : int, optional
+        Minimum size for each dimension. Default is 1.
+    max_size : int, optional
+        Maximum size for each dimension. Default is 100.
+
+    Returns
+    -------
+    tuple[tuple[int, ...], tuple[int, ...]]
+        A tuple containing (shape, chunks) where:
+        - shape is a tuple of positive integers
+        - chunks is a tuple where each element is an integer <= corresponding shape element
+
+    Examples
+    --------
+    >>> shape_and_chunks().example()  # doctest: +SKIP
+    ((5, 3, 8), (2, 3, 4))
+    >>> shape_and_chunks().example()  # doctest: +SKIP
+    ((10, 7), (10, 3))
+    >>> shape_and_chunks(min_dims=2, max_dims=3).example()  # doctest: +SKIP
+    ((4, 6, 2), (2, 3, 1))
+
+    See Also
+    --------
+    :ref:`testing.hypothesis`_
+    """
+    # Generate the shape tuple
+    ndim = draw(st.integers(min_value=min_dims, max_value=max_dims))
+    shape = draw(
+        st.tuples(
+            *[st.integers(min_value=min_size, max_value=max_size) for _ in range(ndim)]
+        )
+    )
+
+    # Generate chunks tuple with each element <= corresponding shape element
+    chunks_elements = []
+    for size in shape:
+        # Each chunk is an integer between 1 and the size of that dimension
+        chunk_element = draw(st.integers(min_value=1, max_value=size))
+        chunks_elements.append(chunk_element)
+
+    chunks = tuple(chunks_elements)
+    return shape, chunks
+
+
 _readable_strings = st.text(
     _readable_characters,
     max_size=5,