Code review: non-optimal way to test filtration generator

1 files changed, 19 insertions, 26 deletions

diff --git a/src/python/test/test_euclidean_witness_complex.py b/src/python/test/test_euclidean_witness_complex.py
index 16ff1ef4..47196a2a 100755
--- a/src/python/test/test_euclidean_witness_complex.py
+++ b/src/python/test/test_euclidean_witness_complex.py
@@ -41,16 +41,15 @@ def test_witness_complex():
     assert landmarks[1] == euclidean_witness_complex.get_point(1)
     assert landmarks[2] == euclidean_witness_complex.get_point(2)
 
-    filtration_generator = simplex_tree.get_filtration()
-    assert(next(filtration_generator) == ([0], 0.0))
-    assert(next(filtration_generator) == ([1], 0.0))
-    assert(next(filtration_generator) == ([0, 1], 0.0))
-    assert(next(filtration_generator) == ([2], 0.0))
-    assert(next(filtration_generator) == ([0, 2], 0.0))
-    assert(next(filtration_generator) == ([1, 2], 0.0))
-    assert(next(filtration_generator) == ([0, 1, 2], 0.0))
-    with pytest.raises(StopIteration):
-        next(filtration_generator)
+    assert list(simplex_tree.get_filtration()) == [
+        ([0], 0.0),
+        ([1], 0.0),
+        ([0, 1], 0.0),
+        ([2], 0.0),
+        ([0, 2], 0.0),
+        ([1, 2], 0.0),
+        ([0, 1, 2], 0.0),
+    ]
 
 
 def test_empty_euclidean_strong_witness_complex():
@@ -80,24 +79,18 @@ def test_strong_witness_complex():
     assert landmarks[1] == euclidean_strong_witness_complex.get_point(1)
     assert landmarks[2] == euclidean_strong_witness_complex.get_point(2)
 
-    filtration_generator = simplex_tree.get_filtration()
-    assert(next(filtration_generator) == ([0], 0.0))
-    assert(next(filtration_generator) == ([1], 0.0))
-    assert(next(filtration_generator) == ([2], 0.0))
-    with pytest.raises(StopIteration):
-        next(filtration_generator)
+    assert list(simplex_tree.get_filtration()) == [([0], 0.0), ([1], 0.0), ([2], 0.0)]
 
     simplex_tree = euclidean_strong_witness_complex.create_simplex_tree(
         max_alpha_square=100.0
     )
 
-    filtration_generator = simplex_tree.get_filtration()
-    assert(next(filtration_generator) == ([0], 0.0))
-    assert(next(filtration_generator) == ([1], 0.0))
-    assert(next(filtration_generator) == ([2], 0.0))
-    assert(next(filtration_generator) == ([1, 2], 15.0))
-    assert(next(filtration_generator) == ([0, 2], 34.0))
-    assert(next(filtration_generator) == ([0, 1], 37.0))
-    assert(next(filtration_generator) == ([0, 1, 2], 37.0))
-    with pytest.raises(StopIteration):
-        next(filtration_generator)
+    assert list(simplex_tree.get_filtration()) == [
+        ([0], 0.0),
+        ([1], 0.0),
+        ([2], 0.0),
+        ([1, 2], 15.0),
+        ([0, 2], 34.0),
+        ([0, 1], 37.0),
+        ([0, 1, 2], 37.0),
+    ]