Merge branch 'alan-turing-institute:main' into main

Author: James Parkhurst

_tests/test_io.py

@@ -158,60 +158,52 @@ def test_command_line_main(tmpdir, test_id):
         assert os.path.exists(filename)
 
 
-def get_cleave_ids():
-    test_id_map = {
-        "alphafold": ["P0A855_cleaved_1to21", "P21170_cleaved_none"],
-        "pdb": ["P0A855_1c5k_cleaved_1to21", "P21170_3nzq_cleaved_none"],
-    }
-    for source, test_id_list in test_id_map.items():
-        for test_id in test_id_list:
-            yield source, test_id
+@pytest.mark.parametrize(
+    "signal_peptides, hydrogens, water, hetatoms, expected_filename_suffix",
+    [
+        (True, True, True, True, "nosignal1to25_nohydrogens_nowater_nohetatm"),
+        (False, True, True, True, "nohydrogens_nowater_nohetatm"),
+        (True, False, True, True, "nosignal1to25_nowater_nohetatm"),
+        (True, True, False, True, "nosignal1to25_nohydrogens_nohetatm"),
+        (True, True, True, False, "nosignal1to25_nohydrogens_nowater"),
+        (
+            False,
+            False,
+            False,
+            False,
+            "unmodified",
+        ),
+    ],
+)
+def test_fetcher_remove_stuff(
+    tmpdir,
+    signal_peptides,
+    hydrogens,
+    water,
+    hetatoms,
+    expected_filename_suffix,
+):
+    fetcher = Fetcher()
+    fetcher.set_directory(str(tmpdir))
 
+    # Fetch the file
+    fetcher.get_file(
+        uniprot_id="P45523", filetype="pdb", filesave=True, db="pdb"
+    )
 
-def remove_suffixes(ids):
-    # Define patterns to match "_none" and "_1to21" at the end of the string
-    patterns = [r"_cleaved_none$", r"_cleaved_1to21$"]
-    # Apply each pattern and remove matching suffixes
-    for pattern in patterns:
-        ids = re.sub(pattern, "", ids)
-    return ids
+    # Remove elements based on the test parameters
+    fetcher.remove(
+        "P45523",
+        signal_peptides=signal_peptides,
+        hydrogens=hydrogens,
+        water=water,
+        hetatoms=hetatoms,
+    )
 
+    # Construct the expected filename
+    expected_filename = os.path.join(
+        str(tmpdir), f"p45523_1q6u_{expected_filename_suffix}.pdb"
+    )
 
-@pytest.mark.parametrize("test_id", get_cleave_ids())
-def test_fetcher_cleave_off_signal_peptides_pdb(tmpdir, test_id):
-    source, pdb_id_signals = test_id
-    pdb_id = remove_suffixes(pdb_id_signals)
-    fetcher = Fetcher(source)
-    fetcher.set_directory(str(tmpdir))
-    if source == "pdb":
-        pdb_id = pdb_id[:-5]
-        fetcher.get_file(
-            uniprot_id=pdb_id, filetype="pdb", filesave=True, db=source
-        )
-        fetcher.cleave_off_signal_peptides(pdb_id)
-    else:
-        fetcher.get_file(
-            uniprot_id=pdb_id, filetype="pdb", filesave=True, db=source
-        )
-        fetcher.cleave_off_signal_peptides(pdb_id)
-    assert os.path.exists(str(tmpdir) + "/" + pdb_id_signals.lower() + ".pdb")
-
-
-@pytest.mark.parametrize("test_id", get_cleave_ids())
-def test_fetcher_cleave_off_signal_peptides_cif(tmpdir, test_id):
-    source, pdb_id_signals = test_id
-    pdb_id = remove_suffixes(pdb_id_signals)
-    fetcher = Fetcher(source)
-    fetcher.set_directory(str(tmpdir))
-    if source == "pdb":
-        pdb_id = pdb_id[:-5]
-        fetcher.get_file(
-            uniprot_id=pdb_id, filetype="cif", filesave=True, db=source
-        )
-        fetcher.cleave_off_signal_peptides(pdb_id)
-    else:
-        fetcher.get_file(
-            uniprot_id=pdb_id, filetype="cif", filesave=True, db=source
-        )
-        fetcher.cleave_off_signal_peptides(pdb_id)
-    assert os.path.exists(str(tmpdir) + "/" + pdb_id_signals.lower() + ".cif")
+    # Assert that the expected file exists
+    assert os.path.exists(expected_filename)