diff --git a/tests/test_custom_return_types.py b/tests/test_custom_return_types.py index 48aeb7c0a..504a9f734 100644 --- a/tests/test_custom_return_types.py +++ b/tests/test_custom_return_types.py @@ -95,6 +95,16 @@ def test_slices(self): slice_return = successors[0:3:2] self.assertEqual([("a", ["c", "b"])], slice_return) + def test_iter(self): + self.dag.add_child(self.node_a, "c", "edge") + successors = rustworkx.bfs_successors(self.dag, 0) + self.assertEqual(list(successors), list(iter(successors))) + + def test_reversed(self): + self.dag.add_child(self.node_a, "c", "edge") + successors = rustworkx.bfs_successors(self.dag, 0) + self.assertEqual(list(successors)[::-1], list(reversed(successors))) + class TestNodeIndicesComparisons(unittest.TestCase): def setUp(self): @@ -174,6 +184,10 @@ def test_slices_negatives(self): self.assertEqual([2, 3], slice_return) self.assertEqual([], indices[-1:-2]) + def test_iter(self): + indices = self.dag.node_indices() + self.assertEqual(list(indices), list(iter(indices))) + def test_reversed(self): indices = self.dag.node_indices() reversed_slice = indices[::-1] @@ -352,6 +366,14 @@ def test_slices(self): slice_return = edges[0:-1] self.assertEqual([0, 1], slice_return) + def test_iter(self): + indices = self.dag.edge_indices() + self.assertEqual(list(indices), list(iter(indices))) + + def test_reversed(self): + indices = self.dag.edge_indices() + self.assertEqual(list(indices)[::-1], list(reversed(indices))) + class TestEdgeListComparisons(unittest.TestCase): def setUp(self): @@ -425,6 +447,14 @@ def test_numpy_conversion(): np.asarray(res, dtype=np.uintp), np.array([[0, 1], [0, 2], [0, 3], [0, 4]]) ) + def test_iter(self): + edges = self.dag.edge_list() + self.assertEqual(list(edges), list(iter(edges))) + + def test_reversed(self): + edges = self.dag.edge_list() + self.assertEqual(list(edges)[::-1], list(reversed(edges))) + class TestWeightedEdgeListComparisons(unittest.TestCase): def setUp(self): @@ -500,6 +530,14 @@ def test_numpy_conversion(self): np.asarray(self.dag.weighted_edge_list()), np.array([(0, 1, "Edgy")], dtype=object) ) + def test_iter(self): + edges = self.dag.weighted_edge_list() + self.assertEqual(list(edges), list(iter(edges))) + + def test_reversed(self): + edges = self.dag.weighted_edge_list() + self.assertEqual(list(edges)[::-1], list(reversed(edges))) + class TestPathMapping(unittest.TestCase): def setUp(self): @@ -1387,6 +1425,12 @@ def test_numpy_conversion(self): # this test assumes the array is 1-dimensional which avoids issues with jagged arrays self.assertTrue(np.asarray(self.chains).shape, (1,)) + def test_iter(self): + self.assertEqual(list(self.chains), list(iter(self.chains))) + + def test_reversed(self): + self.assertEqual(list(self.chains)[::-1], list(reversed(self.chains))) + class TestProductNodeMap(unittest.TestCase): def setUp(self):