Enhance test_diag_indices.py and test_flatten.py.

tests/integration/test_append.py

@@ -15,7 +15,7 @@
 
 import numpy as np
 import pytest
-from utils.utils import check_array_method
+from utils.utils import check_module_function
 
 import cunumeric as num
 
@@ -45,7 +45,7 @@ def test_append(size):
         size_b[axis] = size[axis] + 10
         b = np.random.randint(low=0, high=100, size=size_b)
         print_msg = f"np.append(array({a.shape}), array({b.shape}), {axis})"
-        check_array_method("append", [a, b], {"axis": axis}, print_msg)
+        check_module_function("append", [a, b], {"axis": axis}, print_msg)
 
 
 @pytest.mark.parametrize("size_b", SIZES, ids=str)
@@ -55,7 +55,7 @@ def test_append_axis_none(size_a, size_b):
     a = np.random.randint(low=0, high=100, size=size_a)
     b = np.random.randint(low=0, high=100, size=size_b)
     print_msg = f"np.append(array({a.shape}), array({b.shape}), {axis})"
-    check_array_method("append", [a, b], {"axis": axis}, print_msg)
+    check_module_function("append", [a, b], {"axis": axis}, print_msg)
 
 
 class TestAppendErrors:

tests/integration/test_array_split.py

@@ -17,7 +17,7 @@
 
 import numpy as np
 import pytest
-from utils.utils import check_array_method
+from utils.utils import check_module_function
 
 import cunumeric as num
 
@@ -130,7 +130,7 @@ def test_array_split(size):
             # Check if both impls produce the error
             # for non-viable options
 
-            check_array_method(
+            check_module_function(
                 "array_split", [a, input_opt], {"axis": axis}, print_msg
             )
 

tests/integration/test_block.py

@@ -15,7 +15,7 @@
 
 import numpy as np
 import pytest
-from utils.utils import check_array_method
+from utils.utils import check_module_function
 
 import cunumeric as num
 
@@ -48,7 +48,7 @@ def test(depth, sizes):
         f"{_deepen(depth, b)}])"
     )
     arg = [_deepen(depth, a), _deepen(depth, b)]
-    check_array_method("block", [arg], {}, print_msg, check_type=False)
+    check_module_function("block", [arg], {}, print_msg, check_type=False)
 
 
 class TestBlock:
@@ -61,7 +61,7 @@ def test_block_simple_row_wise(self):
             f"np & cunumeric.block([array({a_2d.shape}), "
             f"array({b_2d.shape})])"
         )
-        check_array_method("block", [arg], {}, print_msg)
+        check_module_function("block", [arg], {}, print_msg)
 
     def test_block_simple_column_wise(self):
         a_2d = np.ones((2, 2))
@@ -72,7 +72,7 @@ def test_block_simple_column_wise(self):
             f"np & cunumeric.block([[array({a_2d.shape})], "
             f"[array({b_2d.shape})]])"
         )
-        check_array_method("block", [arg], {}, print_msg)
+        check_module_function("block", [arg], {}, print_msg)
 
     def test_block_with_1d_arrays_multiple_rows(self):
         a = np.array([1, 2, 3])
@@ -83,7 +83,7 @@ def test_block_with_1d_arrays_multiple_rows(self):
             f"np & cunumeric.block([[array({a.shape}), array({b.shape})], "
             f"[array({a.shape}), array({b.shape})]])"
         )
-        check_array_method("block", [arg], {}, print_msg, check_type=False)
+        check_module_function("block", [arg], {}, print_msg, check_type=False)
 
     def test_block_mixed_1d_and_2d(self):
         a_2d = np.ones((2, 2))
@@ -94,7 +94,7 @@ def test_block_mixed_1d_and_2d(self):
             f"np & cunumeric.block([[array({a_2d.shape})], "
             f"[array({b_1d.shape})]])"
         )
-        check_array_method("block", [arg], {}, print_msg)
+        check_module_function("block", [arg], {}, print_msg)
 
     def test_block_complicated(self):
         one_2d = np.array([[1, 1, 1]])
@@ -113,7 +113,7 @@ def test_block_complicated(self):
         ]
 
         print_msg = "np & cunumeric.block()"
-        check_array_method("block", [arg], {}, print_msg)
+        check_module_function("block", [arg], {}, print_msg)
 
     def test_nested(self):
         one = np.array([1, 1, 1])
@@ -165,7 +165,7 @@ def test_3d(self):
         ]
 
         print_msg = "np & cunumeric.block()"
-        check_array_method("block", [arg], {}, print_msg, check_type=False)
+        check_module_function("block", [arg], {}, print_msg, check_type=False)
 
 
 class TestBlockErrors:

tests/integration/test_diag_indices.py

@@ -20,23 +20,104 @@
 import cunumeric as cn
 
 
+def test_diag_indices_default_ndim():
+    np_res = np.diag_indices(10)
+    cn_res = cn.diag_indices(10)
+    assert np.array_equal(np_res, cn_res)
+
+
 @pytest.mark.parametrize("ndim", range(0, LEGATE_MAX_DIM + 1))
-def test_diag_indices(ndim):
+def test_diag_indices_basic(ndim):
     np_res = np.diag_indices(10, ndim)
     cn_res = cn.diag_indices(10, ndim)
     assert np.array_equal(np_res, cn_res)
 
 
+@pytest.mark.parametrize("n", [0, 0.0, 1, 10.5])
+@pytest.mark.parametrize("ndim", [-4, 0, 1])
+def test_diag_indices(n, ndim):
+    np_res = np.diag_indices(n, ndim)
+    cn_res = cn.diag_indices(n, ndim)
+    assert np.array_equal(np_res, cn_res)
+
+
+class TestDiagIndicesErrors:
+    @pytest.mark.parametrize("n", [-10.5, -1])
+    def test_negative_n(self, n):
+        with pytest.raises(ValueError):
+            cn.diag_indices(n)
+
+    @pytest.mark.xfail
+    @pytest.mark.parametrize("n", [-10.5, -1])
+    def test_negative_n_DIVERGENCE(self, n):
+        # np.diag_indices(-10.5) returns empty 2-D array, dtype=float64
+        # np.diag_indices(-1) returns empty 2-D array, dtype=int32
+        # cn.diag_indices(-10.5) raises ValueError
+        # cn.diag_indices(-1) raises ValueError
+        np_res = np.diag_indices(n)
+        cn_res = cn.diag_indices(n)
+        assert np.array_equal(np_res, cn_res)
+
+    def test_none_n(self):
+        msg = "unsupported operand type"
+        with pytest.raises(TypeError, match=msg):
+            cn.diag_indices(None)
+
+    @pytest.mark.parametrize("ndim", [-1.5, 0.0, 1.5])
+    def test_float_ndim(self, ndim):
+        msg = "can't multiply sequence by non-int of type 'float'"
+        with pytest.raises(TypeError, match=msg):
+            cn.diag_indices(10, ndim)
+
+    def test_none_ndim(self):
+        msg = "can't multiply sequence by non-int of type 'NoneType'"
+        with pytest.raises(TypeError, match=msg):
+            cn.diag_indices(10, None)
+
+
+@pytest.mark.parametrize("size", [(5,), (0,)], ids=str)
 @pytest.mark.parametrize("ndim", range(2, LEGATE_MAX_DIM + 1))
-def test_diag_indices_from(ndim):
-    shape = (5,) * ndim
+def test_diag_indices_from_basic(size, ndim):
+    shape = size * ndim
     a = np.ones(shape, dtype=int)
     a_cn = cn.array(a)
     np_res = np.diag_indices_from(a)
     cn_res = cn.diag_indices_from(a_cn)
     assert np.array_equal(np_res, cn_res)
 
 
+class TestDiagIndicesFromErrors:
+    @pytest.mark.parametrize("size", [(5,), (0,)], ids=str)
+    def test_1d(self, size):
+        a = np.ones(size, dtype=int)
+        msg = "input array must be at least 2-d"
+        with pytest.raises(ValueError, match=msg):
+            cn.diag_indices_from(a)
+
+    def test_0d(self):
+        a = np.array(3)
+        msg = "input array must be at least 2-d"
+        with pytest.raises(ValueError, match=msg):
+            cn.diag_indices_from(a)
+
+    @pytest.mark.parametrize(
+        "size",
+        [
+            (
+                5,
+                2,
+            ),
+            (5, 0, 0),
+        ],
+        ids=str,
+    )
+    def test_unequal_length(self, size):
+        a = np.ones(size, dtype=int)
+        msg = "All dimensions of input must be of equal length"
+        with pytest.raises(ValueError, match=msg):
+            cn.diag_indices_from(a)
+
+
 if __name__ == "__main__":
     import sys
 

tests/integration/test_eye.py

@@ -15,7 +15,7 @@
 
 import numpy as np
 import pytest
-from utils.utils import check_array_method
+from utils.utils import check_module_function
 
 import cunumeric as num
 
@@ -27,27 +27,27 @@
 @pytest.mark.parametrize("M", [N, N + 1, N - 1, N * 10, 0])
 def test_eye(M, k):
     print_msg = f"np & cunumeric.eye({N},{M}, k={k})"
-    check_array_method("eye", [N, M], {"k": k}, print_msg)
+    check_module_function("eye", [N, M], {"k": k}, print_msg)
 
 
 @pytest.mark.parametrize("dtype", [np.int32, np.float64, None], ids=str)
 @pytest.mark.parametrize("k", KS, ids=str)
 def test_square(k, dtype):
     print_msg = f"np & cunumeric.eye({N},k={k},dtype={dtype})"
-    check_array_method("eye", [N], {"k": k, "dtype": dtype}, print_msg)
+    check_module_function("eye", [N], {"k": k, "dtype": dtype}, print_msg)
 
 
 def test_N_zero():
     N = 0
     print_msg = f"np & cunumeric eye({N})"
-    check_array_method("eye", [N], {}, print_msg)
+    check_module_function("eye", [N], {}, print_msg)
 
 
 def test_M_zero():
     N = 5
     M = 0
     print_msg = f"np & cunumeric eye({N},{M})"
-    check_array_method("eye", [N, M], {}, print_msg)
+    check_module_function("eye", [N, M], {}, print_msg)
 
 
 class TestEyeErrors:

tests/integration/test_flatten.py

@@ -15,41 +15,9 @@
 
 import numpy as np
 import pytest
+from utils.utils import check_array_method
 
-from cunumeric.array import convert_to_cunumeric_ndarray
-
-
-def run_test(np_arr, num_arr):
-    # We don't support 'K' yet, which will be supported later
-    test_orders = ["C", "F", "A"]
-    for order in test_orders:
-        b = np_arr.flatten(order)
-        c = num_arr.flatten(order)
-        is_equal = True
-        err_arr = [b, c]
-
-        if len(b) != len(c):
-            is_equal = False
-            err_arr = [b, c]
-        else:
-            for each in zip(b, c):
-                if not np.array_equal(*each):
-                    err_arr = each
-                    is_equal = False
-                    break
-        print_msg = f"np & cunumeric.ndarray({np_arr.shape}).flatten({order}))"
-        assert is_equal, (
-            f"Failed, {print_msg}\n"
-            f"numpy result: {err_arr[0]}, {b.shape}\n"
-            f"cunumeric_result: {err_arr[1]}, {c.shape}\n"
-            f"cunumeric and numpy shows"
-            f" different result\n"
-        )
-        print(
-            f"Passed, {print_msg}, np: ({b.shape}, {b.dtype})"
-            f", cunumeric: ({c.shape}, {c.dtype}"
-        )
-
+import cunumeric as num
 
 DIM = 10
 
@@ -67,11 +35,51 @@ def run_test(np_arr, num_arr):
 
 
 # test ndarray.flatten w/ 1D, 2D and 3D arrays
+@pytest.mark.parametrize("order", ("C", "F", "A"))
 @pytest.mark.parametrize("size", SIZES, ids=str)
-def test_basic(size):
+def test_basic(order, size):
     a = np.random.randint(low=0, high=100, size=size)
-    b = convert_to_cunumeric_ndarray(a)
-    run_test(a, b)
+    print_msg = f"np & cunumeric.ndarray.flatten({order})"
+    check_array_method(a, "flatten", [order], {}, print_msg)
+
+
+class TestNdarrayFlattenErrors:
+    def setup(self):
+        size_a = (1, DIM)
+        anp = np.random.randint(low=0, high=100, size=size_a)
+        self.anum = num.array(anp)
+
+    def test_non_string_order(self):
+        order = 0
+        msg = "order must be str, not int"
+        with pytest.raises(TypeError, match=msg):
+            self.anum.flatten(order)
+
+        order = 1
+        msg = "order must be str, not int"
+        with pytest.raises(TypeError, match=msg):
+            self.anum.flatten(order)
+
+        order = -1
+        msg = "order must be str, not int"
+        with pytest.raises(TypeError, match=msg):
+            self.anum.flatten(order)
+
+        order = 1.0
+        msg = "order must be str, not float"
+        with pytest.raises(TypeError, match=msg):
+            self.anum.flatten(order)
+
+        order = ["C"]
+        msg = "order must be str, not list"
+        with pytest.raises(TypeError, match=msg):
+            self.anum.flatten(order)
+
+    def test_bad_string_order(self):
+        order = "Z"
+        msg = "order must be one of 'C', 'F', 'A', or 'K'"
+        with pytest.raises(ValueError, match=msg):
+            self.anum.flatten(order)
 
 
 if __name__ == "__main__":