1 files changed, 31 insertions, 6 deletions
diff --git a/src/python/test/test_alpha_complex.py b/src/python/test/test_alpha_complex.py
index 712a50b6..3761fe16 100755
--- a/src/python/test/test_alpha_complex.py
+++ b/src/python/test/test_alpha_complex.py
@@ -11,8 +11,13 @@
 from gudhi import AlphaComplex, SimplexTree
 import math
 import numpy as np
-import itertools
 import pytest
+try:
+    # python3
+    from itertools import zip_longest
+except ImportError:
+    # python2
+    from itertools import izip_longest as zip_longest
 
 __author__ = "Vincent Rouvreau"
 __copyright__ = "Copyright (C) 2016 Inria"
@@ -60,8 +65,18 @@ def test_infinite_alpha():
     assert point_list[1] == alpha_complex.get_point(1)
     assert point_list[2] == alpha_complex.get_point(2)
     assert point_list[3] == alpha_complex.get_point(3)
-    assert alpha_complex.get_point(4) == []
-    assert alpha_complex.get_point(125) == []
+    try:
+        alpha_complex.get_point(4) == []
+    except IndexError:
+        pass
+    else:
+        assert False
+    try:
+        alpha_complex.get_point(125) == []
+    except IndexError:
+        pass
+    else:
+        assert False
 
 
 def test_filtered_alpha():
@@ -77,8 +92,18 @@ def test_filtered_alpha():
     assert point_list[1] == filtered_alpha.get_point(1)
     assert point_list[2] == filtered_alpha.get_point(2)
     assert point_list[3] == filtered_alpha.get_point(3)
-    assert filtered_alpha.get_point(4) == []
-    assert filtered_alpha.get_point(125) == []
+    try:
+        filtered_alpha.get_point(4) == []
+    except IndexError:
+        pass
+    else:
+        assert False
+    try:
+        filtered_alpha.get_point(125) == []
+    except IndexError:
+        pass
+    else:
+        assert False
 
     assert simplex_tree.get_filtration() == [
         ([0], 0.0),
@@ -114,6 +139,6 @@ def test_safe_alpha_persistence_comparison():
     diag1 = simplex_tree1.persistence()
     diag2 = simplex_tree2.persistence()
 
-    for (first_p, second_p) in itertools.zip_longest(diag1, diag2):
+    for (first_p, second_p) in zip_longest(diag1, diag2):
         assert first_p[0] == pytest.approx(second_p[0])
         assert first_p[1] == pytest.approx(second_p[1])