try to fix tests (#170)

use some smaller test pdbs for tests

Author: lucidrains

tests/test_af3.py

@@ -63,6 +63,8 @@
     default_extract_atompair_feats_fn
 )
 
+DATA_TEST_PDB_ID = '721p'
+
 def test_atom_ref_pos_to_atompair_inputs():
     atom_ref_pos = torch.randn(16, 3)
     atom_ref_space_uid = torch.ones(16).long()
@@ -1095,7 +1097,7 @@ def test_unresolved_protein_rasa():
 
     # rest of the test
 
-    mmcif_filepath = os.path.join('data', 'test', '7a4d-assembly1.cif')
+    mmcif_filepath = os.path.join('data', 'test', f'{DATA_TEST_PDB_ID}-assembly1.cif')
     pdb_input = PDBInput(mmcif_filepath)
 
     mol_input = pdb_input_to_molecule_input(pdb_input)

tests/test_input.py

@@ -26,6 +26,8 @@
 
 from alphafold3_pytorch.mocks import MockAtomDataset
 
+DATA_TEST_PDB_ID = '7a4d'
+
 # reverse complements
 
 def test_string_reverse_complement():
@@ -187,7 +189,7 @@ def test_atompos_input():
 
 def test_pdbinput_input():
     """Test the PDBInput class, particularly its input transformations for mmCIF files."""
-    filepath = os.path.join("data", "test", "7a4d-assembly1.cif")
+    filepath = os.path.join("data", "test", f"{DATA_TEST_PDB_ID}-assembly1.cif")
     file_id = os.path.splitext(os.path.basename(filepath))[0]
     assert os.path.exists(filepath), f"File {filepath} does not exist."
 

tests/test_trainer.py

@@ -24,6 +24,8 @@
 
 from alphafold3_pytorch.mocks import MockAtomDataset
 
+DATA_TEST_PDB_ID = '209d'
+
 def exists(v):
     return v is not None
 
@@ -135,7 +137,7 @@ def test_trainer_with_mock_atom_input(remove_test_folders):
 @pytest.fixture()
 def populate_mock_pdb_and_remove_test_folders():
     proj_root = Path('.')
-    working_cif_file = proj_root / 'data' / 'test' / '7a4d-assembly1.cif'
+    working_cif_file = proj_root / 'data' / 'test' / f'{DATA_TEST_PDB_ID}-assembly1.cif'
 
     pytest_root_folder = Path('./test-folder')
     data_folder = pytest_root_folder / 'data'

tests/test_weighted_sampling.py

@@ -15,6 +15,8 @@
     create_trainer_from_yaml
 )
 
+DATA_TEST_PDB_ID = '209d'
+
 TEST_FOLDER = Path('./data/test')
 
 INTERFACE_MAPPING_PATH = str(TEST_FOLDER / "interface_cluster_mapping.csv")
@@ -54,7 +56,7 @@ def test_cluster_based_sample(sampler: Sampler):
 @pytest.fixture()
 def populate_mock_pdb_and_remove_test_folders():
     proj_root = Path('.')
-    working_cif_file = proj_root / 'data' / 'test' / '7a4d-assembly1.cif'
+    working_cif_file = proj_root / 'data' / 'test' / f'{DATA_TEST_PDB_ID}-assembly1.cif'
 
     pytest_root_folder = Path('./test-folder')
     data_folder = pytest_root_folder / 'data'