Test xpx.testing

Author: crusaderky

pixi.lock

@@ -315,6 +315,7 @@ checks = [
 exclude = [ # don't report on objects that match any of these regex
     '.*test_at.*',
     '.*test_funcs.*',
+    '.*test_testing.*',
     '.*test_utils.*',
     '.*test_version.*',
     '.*test_vendor.*',

pyproject.toml

@@ -6,30 +6,42 @@
     is_pydata_sparse_namespace,
     is_torch_namespace,
 )
-from ._lib._typing import Array
+from ._lib._typing import Array, ModuleType
 
 __all__ = ["xp_assert_close", "xp_assert_equal"]
 
 
-def _check_shape_dtype(actual: Array, desired: Array) -> None:
+def _check_ns_shape_dtype(
+    actual: Array, desired: Array
+) -> ModuleType:  # numpydoc ignore=RT03
     """
-    Assert that shape and dtype of the two arrays match.
+    Assert that namespace, shape and dtype of the two arrays match.
 
     Parameters
     ----------
     actual : Array
         The array produced by the tested function.
     desired : Array
         The expected array (typically hardcoded).
+
+    Returns
+    -------
+    Arrays namespace.
     """
+    actual_xp = array_namespace(actual)  # Raises on scalars and lists
+    desired_xp = array_namespace(desired)
+
+    msg = f"namespaces do not match: {actual_xp} != f{desired_xp}"
+    assert actual_xp == desired_xp, msg
+
     msg = f"shapes do not match: {actual.shape} != f{desired.shape}"
     assert actual.shape == desired.shape, msg
 
-    msg = f"dtypes do not match: {actual.dtype} != {desired.dtype}".format(
-        actual.dtype, desired.dtype
-    )
+    msg = f"dtypes do not match: {actual.dtype} != {desired.dtype}"
     assert actual.dtype == desired.dtype, msg
 
+    return desired_xp
+
 
 def xp_assert_equal(actual: Array, desired: Array, err_msg: str = "") -> None:
     """
@@ -44,8 +56,7 @@ def xp_assert_equal(actual: Array, desired: Array, err_msg: str = "") -> None:
     err_msg : str, optional
         Error message to display on failure.
     """
-    xp = array_namespace(actual, desired)
-    _check_shape_dtype(actual, desired)
+    xp = _check_ns_shape_dtype(actual, desired)
 
     if is_cupy_namespace(xp):
         xp.testing.assert_array_equal(actual, desired, err_msg=err_msg)
@@ -96,8 +107,7 @@ def xp_assert_close(
     err_msg : str, optional
         Error message to display on failure.
     """
-    xp = array_namespace(actual, desired)
-    _check_shape_dtype(actual, desired)
+    xp = _check_ns_shape_dtype(actual, desired)
 
     floating = xp.isdtype(actual.dtype, ("real floating", "complex floating"))
     if rtol is None and floating:

src/array_api_extra/testing.py

@@ -0,0 +1,68 @@
+import numpy as np
+import pytest
+
+from array_api_extra.testing import xp_assert_close, xp_assert_equal
+
+from .conftest import Library
+
+# mypy: disable-error-code=no-any-decorated
+# pyright: reportUnknownParameterType=false,reportMissingParameterType=false
+
+
+@pytest.mark.parametrize(
+    "func",
+    [
+        xp_assert_equal,
+        pytest.param(
+            xp_assert_close,
+            marks=pytest.mark.skip_xp_backend(Library.SPARSE, reason="no isdtype"),
+        ),
+    ],
+)
+def test_assert_close_equal_basic(xp, func):
+    func(xp.asarray(0), xp.asarray(0))
+    func(xp.asarray([1, 2]), xp.asarray([1, 2]))
+
+    with pytest.raises(AssertionError, match="shapes do not match"):
+        func(xp.asarray([0]), xp.asarray([[0]]))
+
+    with pytest.raises(AssertionError, match="dtypes do not match"):
+        func(xp.asarray(0, dtype=xp.float32), xp.asarray(0, dtype=xp.float64))
+
+    with pytest.raises(AssertionError):
+        func(xp.asarray([1, 2]), xp.asarray([1, 3]))
+
+    with pytest.raises(AssertionError, match="hello"):
+        func(xp.asarray([1, 2]), xp.asarray([1, 3]), err_msg="hello")
+
+
+@pytest.mark.skip_xp_backend(Library.NUMPY)
+@pytest.mark.skip_xp_backend(Library.NUMPY_READONLY)
+@pytest.mark.parametrize(
+    "func",
+    [
+        xp_assert_equal,
+        pytest.param(
+            xp_assert_close,
+            marks=pytest.mark.skip_xp_backend(Library.SPARSE, reason="no isdtype"),
+        ),
+    ],
+)
+def test_assert_close_equal_namespace(xp, func):
+    with pytest.raises(AssertionError):
+        func(xp.asarray(0), np.asarray(0))
+    with pytest.raises(TypeError):
+        func(xp.asarray(0), 0)
+    with pytest.raises(TypeError):
+        func(xp.asarray([0]), [0])
+
+
+@pytest.mark.skip_xp_backend(Library.SPARSE, reason="no isdtype")
+def test_assert_close_tolerance(xp):
+    xp_assert_close(xp.asarray([100.0]), xp.asarray([102.0]), rtol=0.03)
+    with pytest.raises(AssertionError):
+        xp_assert_close(xp.asarray([100.0]), xp.asarray([102.0]), rtol=0.01)
+
+    xp_assert_close(xp.asarray([100.0]), xp.asarray([102.0]), atol=3)
+    with pytest.raises(AssertionError):
+        xp_assert_close(xp.asarray([100.0]), xp.asarray([102.0]), atol=1)