Fix pytest

MolecularAI · Jun 3, 2024 · 7daa702 · 7daa702
1 parent f109429
commit 7daa702
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/tests/test_cli.py b/tests/test_cli.py
@@ -419,8 +419,12 @@ def test_download_public_data(tmpdir, mocker, add_cli_arguments):
 
     download_main()
 
-    filenames = glob.glob(str(tmpdir / "*.hdf5"))
-    assert len(filenames) == 4
+    filenames = glob.glob(str(tmpdir / "*.onnx"))
+    assert len(filenames) == 3
+
+    filenames = glob.glob(str(tmpdir / "*.onnx"))
+    assert len(filenames) == 1
+
     for filename in filenames:
         with open(filename, "r") as fileobj:
             assert fileobj.read() == "abcdef"