Update test_flat.py

Author: vlad-perevezentsev

dpnp/tests/test_flat.py

@@ -1,34 +1,53 @@
-import numpy
+import numpy as np
 import pytest
-
-import dpnp as inp
-
-
-@pytest.mark.parametrize("type", [numpy.int64], ids=["int64"])
-def test_flat(type):
-    a = numpy.array([1, 0, 2, -3, -1, 2, 21, -9])
-    ia = inp.array(a)
-
-    result = ia.flat[0]
-    expected = a.flat[0]
-    numpy.testing.assert_array_equal(expected, result)
-
-
-@pytest.mark.parametrize("type", [numpy.int64], ids=["int64"])
-def test_flat2(type):
-    a = numpy.arange(1, 7).reshape(2, 3)
-    ia = inp.array(a)
-
-    result = ia.flat[3]
-    expected = a.flat[3]
-    numpy.testing.assert_array_equal(expected, result)
-
-
-@pytest.mark.parametrize("type", [numpy.int64], ids=["int64"])
-def test_flat3(type):
-    a = numpy.arange(1, 7).reshape(2, 3).T
-    ia = inp.array(a)
-
-    result = ia.flat[3]
-    expected = a.flat[3]
-    numpy.testing.assert_array_equal(expected, result)
+from numpy.testing import assert_array_equal
+
+import dpnp
+
+
+class TestFlatiter:
+    @pytest.mark.parametrize(
+        "a, index",
+        [
+            (np.array([1, 0, 2, -3, -1, 2, 21, -9]), 0),
+            (np.arange(1, 7).reshape(2, 3), 3),
+            (np.arange(1, 7).reshape(2, 3).T, 3),
+        ],
+        ids=["1D array", "2D array", "2D.T array"],
+    )
+    def test_flat_getitem(self, a, index):
+        a_dp = dpnp.array(a)
+        result = a_dp.flat[index]
+        expected = a.flat[index]
+        assert_array_equal(expected, result)
+
+    def test_flat_iteration(self):
+        a = np.array([[1, 2], [3, 4]])
+        a_dp = dpnp.array(a)
+        for dp_val, np_val in zip(a_dp.flat, a.flat):
+            assert dp_val == np_val
+
+    def test_init_error(self):
+        with pytest.raises(TypeError):
+            _ = dpnp.flatiter([1, 2, 3])
+
+    def test_flat_key_error(self):
+        a_dp = dpnp.array(42)
+        with pytest.raises(KeyError):
+            _ = a_dp.flat[1]
+
+    def test_flat_invalid_key(self):
+        a_dp = dpnp.array([1, 2, 3])
+        flat = dpnp.flatiter(a_dp)
+        # check __getitem__
+        with pytest.raises(TypeError):
+            _ = flat["invalid"]
+        # check __setitem__
+        with pytest.raises(TypeError):
+            flat["invalid"] = 42
+
+    def test_flat_out_of_bounds(self):
+        a_dp = dpnp.array([1, 2, 3])
+        flat = dpnp.flatiter(a_dp)
+        with pytest.raises(IndexError):
+            _ = flat[10]