Include test coverage for BindingAffinity.

aion-labs · Aug 7, 2024 · bb4d7f5 · bb4d7f5
1 parent de242cd
commit bb4d7f5
Showing 1 changed file with 11 additions and 1 deletion.
diff --git a/tests/basic_test.py b/tests/basic_test.py
@@ -21,6 +21,7 @@
     FoldswitchProteinsTableS1C,
     CodNas91,
     Weizmann3CA,
+    BindingAffinity,
 )
 
 # List of dataset classes
@@ -44,10 +45,19 @@
     FoldswitchProteinsTableS1C,
     CodNas91,
     Weizmann3CA,
+    BindingAffinity,
 ]
 
+# List of dataset classes that the cache test should not be performed on
+dont_do_cache_test = [
+    BindingAffinity,
+]
+no_cache_datasets = [
+    dataset for dataset in datasets if dataset not in dont_do_cache_test
+]
 
-@pytest.mark.parametrize("dataset_cls", datasets)
+
+@pytest.mark.parametrize("dataset_cls", no_cache_datasets)
 @patch("pathlib.Path.mkdir")
 @patch("polars.DataFrame.write_parquet")
 @patch("pathlib.Path.exists")