Catch warnings in wasserstein test

1 files changed, 5 insertions, 4 deletions

diff --git a/src/python/test/test_wasserstein_distance.py b/src/python/test/test_wasserstein_distance.py
index 3a004d77..a76b6ce7 100755
--- a/src/python/test/test_wasserstein_distance.py
+++ b/src/python/test/test_wasserstein_distance.py
@@ -90,10 +90,11 @@ def test_get_essential_parts():
 
 
 def test_warn_infty():
-    assert _warn_infty(matching=False)==np.inf
-    c, m = _warn_infty(matching=True)
-    assert (c == np.inf)
-    assert (m is None)
+    with pytest.warns(UserWarning):
+        assert _warn_infty(matching=False)==np.inf
+        c, m = _warn_infty(matching=True)
+        assert (c == np.inf)
+        assert (m is None)
 
 
 def _basic_wasserstein(wasserstein_distance, delta, test_infinity=True, test_matching=True):