API: testing: relax condition that desired match xp in assertions

Author: lucascolley

src/array_api_extra/testing.py

@@ -598,22 +598,23 @@ def _check_ns_shape_dtype(
     np = _require_numpy()
 
     actual_xp = array_namespace(actual)  # Raises on Python scalars and lists
-    desired_xp = array_namespace(desired)
 
     if xp is not None:
         _msg = (
-            "Namespace of desired array does not match the `xp` argument.\n"
-            f"Desired array's namespace: {desired_xp.__name__}\n"
+            "Namespace of actual array does not match the `xp` argument.\n"
+            f"Actual array's namespace: {actual_xp.__name__}\n"
             f"Expected namespace: {xp.__name__}."
         )
-        assert desired_xp == xp, _msg
-
-    _msg = (
-        "Namespaces of actual and desired arrays do not match.\n"
-        f"Actual: {actual_xp.__name__}\n"
-        f"Desired: {desired_xp.__name__}."
-    )
-    assert actual_xp == desired_xp, _msg
+        assert actual_xp == xp, _msg
+        desired_xp = xp
+    else:
+        desired_xp = array_namespace(desired)
+        _msg = (
+            "Namespaces of actual and desired arrays do not match.\n"
+            f"Actual: {actual_xp.__name__}\n"
+            f"Desired: {desired_xp.__name__}."
+        )
+        assert actual_xp == desired_xp, _msg
 
     if is_numpy_namespace(actual_xp) and check_scalar:
         # only NumPy distinguishes between scalars and arrays; we do if check_scalar.
@@ -653,7 +654,7 @@ def _check_ns_shape_dtype(
     if check_dtype:
         msg = f"dtypes do not match: {actual.dtype} != {desired.dtype}"
         assert actual.dtype == desired.dtype, msg
-    desired = desired_xp.broadcast_to(desired, actual_shape)
+    desired = desired_xp.broadcast_to(desired_xp.asarray(desired), actual_shape)
     return actual, desired, desired_xp, np
 
 

tests/test_testing.py

@@ -81,10 +81,11 @@ def test_namespace(self, xp: ModuleType, func: Callable[..., None]):
             func(xp.asarray(0), 0)
         with pytest.raises(TypeError, match="list is not a supported array type"):
             func(xp.asarray([0]), [0])
+        func(xp.asarray(0), xp.asarray(1 if func is assert_less else 0), xp=xp)
         with (
             pytest.raises(
                 AssertionError,
-                match="Namespace of desired array does not match the `xp` argument",
+                match="Namespace of actual array does not match the `xp` argument",
             ),
         ):
             func(xp.asarray(0), xp.asarray(0), xp=np)