From 4a87e2eaa3fe5a758a87848dccc6b1dabc7c19f0 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?C=C3=A9dric=20Bouysset?= <cedric@bouysset.net>
Date: Sun, 29 Dec 2024 16:20:18 +0100
Subject: [PATCH] tests: add tests for water-bridge from iterable

---
 tests/conftest.py          | 12 ++++++++++--
 tests/test_fingerprint.py  | 20 ++++++++++++++++----
 tests/test_interactions.py | 18 ++++++++++++++----
 3 files changed, 40 insertions(+), 10 deletions(-)

diff --git a/tests/conftest.py b/tests/conftest.py
index 3f1ca85..1c714cb 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -178,19 +178,27 @@ def water_u():
 
 
 @pytest.fixture(scope="session")
-def water_params(water_u):
+def water_atomgroups(water_u):
     ligand = water_u.select_atoms("resname QNB")
     protein = water_u.select_atoms(
         "protein and byres around 4 group ligand", ligand=ligand
     )
     water = water_u.select_atoms(
-        "resname TIP3 and byres around 6 (group ligand or group pocket)",
+        "resname TIP3 and byres around 4 (group ligand or group pocket)",
         ligand=ligand,
         pocket=protein,
     )
     return ligand, protein, water
 
 
+@pytest.fixture(scope="session")
+def water_mols(water_atomgroups):
+    lig_mol = Molecule.from_mda(water_atomgroups[0])
+    prot_mol = Molecule.from_mda(water_atomgroups[1])
+    water_mol = Molecule.from_mda(water_atomgroups[2])
+    return lig_mol, prot_mol, water_mol
+
+
 class BaseTestMixinRDKitMol:
     def test_init(self, mol):
         assert isinstance(mol, Chem.Mol)
diff --git a/tests/test_fingerprint.py b/tests/test_fingerprint.py
index 9d0ce6e..c9cf86e 100644
--- a/tests/test_fingerprint.py
+++ b/tests/test_fingerprint.py
@@ -345,8 +345,8 @@ def test_water_bridge_instance_without_params_raises_error(self):
         ):
             Fingerprint(["WaterBridge"])
 
-    def test_mix_water_bridge_and_other_interactions(self, water_u, water_params):
-        ligand, protein, water = water_params
+    def test_mix_water_bridge_and_other_interactions(self, water_u, water_atomgroups):
+        ligand, protein, water = water_atomgroups
         fp = Fingerprint(
             ["HBDonor", "WaterBridge"], parameters={"WaterBridge": {"water": water}}
         )
@@ -355,8 +355,20 @@ def test_mix_water_bridge_and_other_interactions(self, water_u, water_params):
         assert "WaterBridge" in fp.ifp[0]["QNB1.X", "TRP400.X"]
         assert "HBDonor" in fp.ifp[0]["QNB1.X", "ASN404.X"]
 
-    def test_water_bridge_updates_cache_size(self, water_u, water_params, monkeypatch):
-        ligand, protein, water = water_params
+    def test_water_bridge_run_iter(self, water_mols):
+        ligand, protein, water = water_mols
+        fp = Fingerprint(
+            ["HBDonor", "WaterBridge"], parameters={"WaterBridge": {"water": water}}
+        )
+        fp.run_from_iterable([ligand], protein)
+
+        assert "WaterBridge" in fp.ifp[0]["QNB1.X", "TRP400.X"]
+        assert "HBDonor" in fp.ifp[0]["QNB1.X", "ASN404.X"]
+
+    def test_water_bridge_updates_cache_size(
+        self, water_u, water_atomgroups, monkeypatch
+    ):
+        ligand, protein, water = water_atomgroups
         set_converter_cache_size(2)
         mocked = Mock(wraps=set_converter_cache_size)
         monkeypatch.setattr(
diff --git a/tests/test_interactions.py b/tests/test_interactions.py
index 6b71dde..79cfbd0 100644
--- a/tests/test_interactions.py
+++ b/tests/test_interactions.py
@@ -357,16 +357,16 @@ class TestBridgedInteractions:
             ({"order": 1, "min_order": 2}, "min_order cannot be greater than order"),
         ],
     )
-    def test_water_bridge_validation(self, water_params, kwargs, match):
-        *_, water = water_params
+    def test_water_bridge_validation(self, water_atomgroups, kwargs, match):
+        *_, water = water_atomgroups
         with pytest.raises(ValueError, match=match):
             Fingerprint(
                 ["WaterBridge"],
                 parameters={"WaterBridge": {"water": water, **kwargs}},
             )
 
-    def test_direct_water_bridge(self, water_u, water_params):
-        ligand, protein, water = water_params
+    def test_direct_water_bridge(self, water_u, water_atomgroups):
+        ligand, protein, water = water_atomgroups
         fp = Fingerprint(["WaterBridge"], parameters={"WaterBridge": {"water": water}})
         fp.run(water_u.trajectory[:1], ligand, protein)
         int_data = next(fp.ifp[0].interactions())
@@ -395,3 +395,13 @@ def test_higher_order_water_bridge(self, water_u, kwargs, num_expected):
         assert len(all_int_data) == num_expected
         int_data = all_int_data[-1]
         assert "distance_TIP383.X_TIP317.X" in int_data.metadata
+
+    def test_run_iter_water_bridge(self, water_mols):
+        ligand, protein, water = water_mols
+        fp = Fingerprint(["WaterBridge"], parameters={"WaterBridge": {"water": water}})
+        # mimick multiple poses
+        fp.run_from_iterable([ligand, ligand], protein)
+        int_data = next(fp.ifp[1].interactions())
+
+        assert int_data.interaction == "WaterBridge"
+        assert str(int_data.protein) == "TRP400.X"