1 files changed, 20 insertions, 35 deletions
diff --git a/debian/tests/upstream.py b/debian/tests/upstream.py
index feadfbd9..defa2da5 100644
--- a/debian/tests/upstream.py
+++ b/debian/tests/upstream.py
@@ -1,42 +1,27 @@
-import importlib
-import os
 import sys
+import itertools
+import pytest
 
-blacklist = {"test_wasserstein_distance", "test_wasserstein_barycenter", "test_knn", "test_subsampling", "test_representations", "test_dtm"}
+blacklist = ["test_knn.py", # Requires pykeops
+             "test_representations.py", # Look into this
+             "test_wasserstein_distance.py::test_wasserstein_distance_grad", # Requires PyTorch.
+             "test_remote_datasets.py" # Requires Internet access, and downloads DFSG-dubious data.
+            ]
 
-testdir = sys.argv[1]
-sys.path.append(testdir)
+print("Blacklist: %s" %(str(blacklist)))
+
+# These should be completely ignored, not even loaded (their imports fail).
+ignorelist = ["test_diff.py", # Requires TF.
+              "test_dtm.py", # Requires hnswlib.
+              "test_wasserstein_with_tensors.py" # Requires PyTorch and TF
+            ]
 
-results = []
-with os.scandir(testdir) as it:
-    for entry in it:
-        if entry.is_file() and entry.name.startswith("test_") and entry.name.endswith(".py"):
-            name = entry.name[:-3]
+print("Ignorelist:", ignorelist)
 
-            if name in blacklist:
-                print("Skipping tests in %s due to blacklist." %(name))
-            else:
-                print("Running tests from %s." %(name))
-            
-                module = importlib.import_module(name)
+testdir = sys.argv[1]
 
-                tests = [f for f in dir(module) if str(f).startswith("test_")]
-                for t in tests:
-                    func = getattr(module, t)
-                    if callable(func):
-                        print(" ", t)
-                        ok = True
-                        try:
-                            func()
-                        except AssertionError:
-                            ok = False
-                        if ok:
-                            print("   OK!")
-                        else:
-                            print("   FAIL!")
-                        results.append(ok)
+deselect_args = list(itertools.chain.from_iterable([["--deselect", "%s/%s" %(testdir, exclude)] for exclude in blacklist]))
+ignore_args = list(itertools.chain.from_iterable([["--ignore-glob", "%s/%s" %(testdir, exclude)] for exclude in ignorelist]))
 
-if all(results):
-    exit(0)
-else:
-    exit(1)
+exitcode = pytest.main([testdir] + deselect_args + ignore_args)
+exit(exitcode)