prepare tests for ndim (#85)

Author: flying-sheep

tests/test_stats.py

@@ -23,16 +23,18 @@
 
     Array: TypeAlias = CpuArray | GpuArray | DiskArray | types.CSDataset | types.DaskArray
 
-    DTypeIn = type[np.float32 | np.float64 | np.int32 | np.bool]
-    DTypeOut = type[np.float32 | np.float64 | np.int64]
+    DTypeIn = np.float32 | np.float64 | np.int32 | np.bool
+    DTypeOut = np.float32 | np.float64 | np.int64
+
+    NdAndAx: TypeAlias = tuple[Literal[2], Literal[0, 1, None]]
 
     class BenchFun(Protocol):  # noqa: D101
         def __call__(  # noqa: D102
             self,
             arr: CpuArray,
             *,
             axis: Literal[0, 1, None] = None,
-            dtype: DTypeOut | None = None,
+            dtype: type[DTypeOut] | None = None,
         ) -> NDArray[Any] | np.number[Any] | types.DaskArray: ...
 
 
@@ -44,29 +46,53 @@ def __call__(  # noqa: D102
 ATS_CUPY_SPARSE = {at for at in SUPPORTED_TYPES if "cupyx.scipy" in str(at)}
 
 
-@pytest.fixture(scope="session", params=[0, 1, None])
-def axis(request: pytest.FixtureRequest) -> Literal[0, 1, None]:
-    return cast("Literal[0, 1, None]", request.param)
+@pytest.fixture(
+    scope="session",
+    params=[
+        pytest.param((2, None), id="2d-all"),
+        pytest.param((2, 0), id="2d-ax0"),
+        pytest.param((2, 1), id="2d-ax1"),
+    ],
+)
+def ndim_and_axis(request: pytest.FixtureRequest) -> NdAndAx:
+    return cast("NdAndAx", request.param)
+
+
+@pytest.fixture
+def ndim(ndim_and_axis: NdAndAx) -> Literal[2]:
+    return ndim_and_axis[0]
+
+
+@pytest.fixture(scope="session")
+def axis(ndim_and_axis: NdAndAx) -> Literal[0, 1, None]:
+    return ndim_and_axis[1]
 
 
 @pytest.fixture(scope="session", params=[np.float32, np.float64, np.int32, np.bool])
-def dtype_in(request: pytest.FixtureRequest) -> DTypeIn:
-    return cast("DTypeIn", request.param)
+def dtype_in(request: pytest.FixtureRequest) -> type[DTypeIn]:
+    return cast("type[DTypeIn]", request.param)
 
 
 @pytest.fixture(scope="session", params=[np.float32, np.float64, None])
-def dtype_arg(request: pytest.FixtureRequest) -> DTypeOut | None:
-    return cast("DTypeOut | None", request.param)
+def dtype_arg(request: pytest.FixtureRequest) -> type[DTypeOut] | None:
+    return cast("type[DTypeOut] | None", request.param)
+
+
+@pytest.fixture
+def np_arr(dtype_in: type[DTypeIn]) -> NDArray[DTypeIn]:
+    np_arr = cast("NDArray[DTypeIn]", np.array([[1, 2, 3], [4, 5, 6]], dtype=dtype_in))
+    np_arr.flags.writeable = False
+    return np_arr
 
 
 @pytest.mark.array_type(skip=ATS_SPARSE_DS)
 def test_sum(
     array_type: ArrayType[Array],
-    dtype_in: DTypeIn,
-    dtype_arg: DTypeOut | None,
+    dtype_in: type[DTypeIn],
+    dtype_arg: type[DTypeOut] | None,
     axis: Literal[0, 1, None],
+    np_arr: NDArray[DTypeIn],
 ) -> None:
-    np_arr = np.array([[1, 2, 3], [4, 5, 6]], dtype=dtype_in)
     if array_type in ATS_CUPY_SPARSE and np_arr.dtype.kind != "f":
         pytest.skip("CuPy sparse matrices only support floats")
     arr = array_type(np_arr.copy())
@@ -104,21 +130,20 @@ def test_sum(
 
 
 @pytest.mark.array_type(skip=ATS_SPARSE_DS)
-@pytest.mark.parametrize(("axis", "expected"), [(None, 3.5), (0, [2.5, 3.5, 4.5]), (1, [2.0, 5.0])])
 def test_mean(
-    array_type: ArrayType[Array], axis: Literal[0, 1, None], expected: float | list[float]
+    array_type: ArrayType[Array], axis: Literal[0, 1, None], np_arr: NDArray[DTypeIn]
 ) -> None:
-    np_arr = np.array([[1, 2, 3], [4, 5, 6]])
     if array_type in ATS_CUPY_SPARSE and np_arr.dtype.kind != "f":
         pytest.skip("CuPy sparse matrices only support floats")
-    np.testing.assert_array_equal(np.mean(np_arr, axis=axis), expected)
-
     arr = array_type(np_arr)
+
     result = stats.mean(arr, axis=axis)  # type: ignore[arg-type]  # https://github.com/python/mypy/issues/16777
     if isinstance(result, types.DaskArray):
         result = result.compute()
     if isinstance(result, types.CupyArray | types.CupyCSMatrix):
         result = result.get()
+
+    expected = np.mean(np_arr, axis=axis)  # type: ignore[arg-type]
     np.testing.assert_array_equal(result, expected)