refactor: clean up test data handling (#1815)

* import pooch_cache from conftest * remove unused files from pooch cache * move zenodo cache to conftest * remove duplicate code
2026-06-04 14:14:22 +08:00 · 2026-01-29 07:08:04 -08:00
parent f7561cda22
commit b93807dd10
7 changed files with 63 additions and 123 deletions
--- a/src/openfe/tests/conftest.py
+++ b/src/openfe/tests/conftest.py
@@ -12,6 +12,7 @@ import mdtraj
 import numpy as np
 import openmm
 import pandas as pd
+import pooch
 import pytest
 from gufe import AtomMapper, LigandAtomMapping, ProteinComponent, SmallMoleculeComponent
 from openff.toolkit import ForceField
@@ -26,6 +27,8 @@ from openfe.protocols.openmm_rfe._rfe_utils.relative import HybridTopologyFactor
 from openfe.protocols.openmm_utils.serialization import deserialize
 from openfe.tests.protocols.openmm_rfe.helpers import make_htf

+POOCH_CACHE = pooch.os_cache("openfe")
+

 class SlowTests:
    """Plugin for handling fixtures that skips slow tests
--- a/src/openfe/tests/protocols/conftest.py
+++ b/src/openfe/tests/protocols/conftest.py
@@ -1,9 +1,11 @@
 # This code is part of OpenFE and is licensed under the MIT license.
 # For details, see https://github.com/OpenFreeEnergy/openfe
 import gzip
+import pathlib
 from importlib import resources
 from typing import Optional

+import MDAnalysis as mda
 import openmm
 import pooch
 import pytest
@@ -16,6 +18,8 @@ from rdkit.Geometry import Point3D

 import openfe

+from ..conftest import POOCH_CACHE
+

@pytest.fixture
 def available_platforms() -> set[str]:
@@ -280,15 +284,51 @@ def septop_json() -> str:
        return f.read().decode()  # type: ignore


+zenodo_industry_benchmarks_data = pooch.create(
+    path=POOCH_CACHE,
+    base_url="doi:10.5281/zenodo.15212342",
+    registry={
+        "industry_benchmark_systems.zip": "sha256:2bb5eee36e29b718b96bf6e9350e0b9957a592f6c289f77330cbb6f4311a07bd"
+    },
+)
+
+
+@pytest.fixture
+def industry_benchmark_files():
+    zenodo_industry_benchmarks_data.fetch("industry_benchmark_systems.zip", processor=pooch.Unzip())
+    cache_dir = pathlib.Path(
+        POOCH_CACHE / "industry_benchmark_systems.zip.unzip/industry_benchmark_systems"
+    )
+    return cache_dir
+
+
+zenodo_restraint_data = pooch.create(
+    path=POOCH_CACHE,
+    base_url="doi:10.5281/zenodo.15212342",
+    registry={
+        "t4_lysozyme_trajectory.zip": "sha256:e985d055db25b5468491e169948f641833a5fbb67a23dbb0a00b57fb7c0e59c8"
+    },
+)
+
+
+@pytest.fixture
+def t4_lysozyme_trajectory_universe():
+    zenodo_restraint_data.fetch("t4_lysozyme_trajectory.zip", processor=pooch.Unzip())
+    cache_dir = pathlib.Path(
+        POOCH_CACHE / "t4_lysozyme_trajectory.zip.unzip/t4_lysozyme_trajectory"
+    )
+    universe = mda.Universe(
+        str(cache_dir / "t4_toluene_complex.pdb"),
+        str(cache_dir / "t4_toluene_complex.xtc"),
+    )
+    return universe
+
+
 RFE_OUTPUT = pooch.create(
-    path=pooch.os_cache("openfe_analysis"),
+    path=POOCH_CACHE,
    base_url="doi:10.6084/m9.figshare.24101655",
    registry={
-        "checkpoint.nc": "5af398cb14340fddf7492114998b244424b6c3f4514b2e07e4bd411484c08464",
-        "db.json": "b671f9eb4daf9853f3e1645f9fd7c18150fd2a9bf17c18f23c5cf0c9fd5ca5b3",
-        "hybrid_system.pdb": "07203679cb14b840b36e4320484df2360f45e323faadb02d6eacac244fddd517",
        "simulation.nc": "92361a0864d4359a75399470135f56642b72c605069a4c33dbc4be6f91f28b31",
-        "simulation_real_time_analysis.yaml": "65706002f371fafba96037f29b054fd7e050e442915205df88567f48f5e5e1cf",
    },
 )

--- a/src/openfe/tests/protocols/restraints/test_geometry_boresch.py
+++ b/src/openfe/tests/protocols/restraints/test_geometry_boresch.py
@@ -14,7 +14,7 @@ from openfe.protocols.restraint_utils.geometry.boresch.geometry import (
    find_boresch_restraint,
 )

-from ...conftest import HAS_INTERNET
+from ...conftest import HAS_INTERNET, POOCH_CACHE


@pytest.fixture()
@@ -235,25 +235,6 @@ def test_get_boresch_restraint_dssp(eg5_protein_ligand_universe, eg5_ligands):
    assert -0.02396901 == pytest.approx(restraint_geometry.phi_C0.to("radians").m)


-POOCH_CACHE = pooch.os_cache("openfe")
-zenodo_restraint_data = pooch.create(
-    path=POOCH_CACHE,
-    base_url="doi:10.5281/zenodo.15212342",
-    registry={
-        "industry_benchmark_systems.zip": "sha256:2bb5eee36e29b718b96bf6e9350e0b9957a592f6c289f77330cbb6f4311a07bd"
-    },
-)
-
-
-@pytest.fixture
-def industry_benchmark_files():
-    zenodo_restraint_data.fetch("industry_benchmark_systems.zip", processor=pooch.Unzip())
-    cache_dir = pathlib.Path(
-        pooch.os_cache("openfe") / "industry_benchmark_systems.zip.unzip/industry_benchmark_systems"
-    )
-    return cache_dir
-
-
@pytest.mark.skipif(
    not os.path.exists(POOCH_CACHE) and not HAS_INTERNET,
    reason="Internet seems to be unavailable and test data is not cached locally.",
--- a/src/openfe/tests/protocols/restraints/test_geometry_boresch_host.py
+++ b/src/openfe/tests/protocols/restraints/test_geometry_boresch_host.py
@@ -2,7 +2,6 @@
 # For details, see https://github.com/OpenFreeEnergy/openfe

 import os
-import pathlib

 import MDAnalysis as mda
 import numpy as np
@@ -25,31 +24,7 @@ from openfe.protocols.restraint_utils.geometry.utils import (
    is_collinear,
 )

-from ...conftest import HAS_INTERNET
-
-POOCH_CACHE = pooch.os_cache("openfe")
-zenodo_restraint_data = pooch.create(
-    path=POOCH_CACHE,
-    base_url="doi:10.5281/zenodo.15212342",
-    registry={
-        "t4_lysozyme_trajectory.zip": "sha256:e985d055db25b5468491e169948f641833a5fbb67a23dbb0a00b57fb7c0e59c8"
-    },
-)
-
-
-@pytest.fixture(scope="module")
-def t4_lysozyme_trajectory_universe():
-    zenodo_restraint_data.fetch("t4_lysozyme_trajectory.zip", processor=pooch.Unzip())
-    cache_dir = pathlib.Path(
-        pooch.os_cache("openfe") / "t4_lysozyme_trajectory.zip.unzip/t4_lysozyme_trajectory"
-    )
-    universe = mda.Universe(
-        str(cache_dir / "t4_toluene_complex.pdb"),
-        str(cache_dir / "t4_toluene_complex.xtc"),
-    )
-    # guess bonds for the protein atoms
-    universe.select_atoms("protein").guess_bonds()
-    return universe
+from ...conftest import HAS_INTERNET, POOCH_CACHE


@pytest.fixture
@@ -358,7 +333,10 @@ class TestFindAnchorBondedTrajectory(TestFindAnchorMulti):

    @pytest.fixture(scope="class")
    def universe(self, t4_lysozyme_trajectory_universe):
-        return t4_lysozyme_trajectory_universe
+        universe = t4_lysozyme_trajectory_universe
+        # guess bonds for the protein atoms
+        universe.select_atoms("protein").guess_bonds()
+        return universe

    @pytest.fixture(scope="class")
    def host_anchor(self, universe):
--- a/src/openfe/tests/protocols/restraints/test_geometry_utils.py
+++ b/src/openfe/tests/protocols/restraints/test_geometry_utils.py
@@ -3,12 +3,10 @@

 import itertools
 import os
-import pathlib
 from importlib import resources

 import MDAnalysis as mda
 import numpy as np
-import pooch
 import pytest
 from openff.units import unit
 from rdkit import Chem
@@ -32,7 +30,7 @@ from openfe.protocols.restraint_utils.geometry.utils import (
    stable_secondary_structure_selection,
 )

-from ...conftest import HAS_INTERNET
+from ...conftest import HAS_INTERNET, POOCH_CACHE


@pytest.fixture(scope="module")
@@ -49,29 +47,6 @@ def eg5_protein_ligand_universe(eg5_protein_pdb, eg5_ligands):
    return mda.Merge(protein.atoms, lig.atoms)


-POOCH_CACHE = pooch.os_cache("openfe")
-zenodo_restraint_data = pooch.create(
-    path=POOCH_CACHE,
-    base_url="doi:10.5281/zenodo.15212342",
-    registry={
-        "t4_lysozyme_trajectory.zip": "sha256:e985d055db25b5468491e169948f641833a5fbb67a23dbb0a00b57fb7c0e59c8"
-    },
-)
-
-
-@pytest.fixture
-def t4_lysozyme_trajectory_universe():
-    zenodo_restraint_data.fetch("t4_lysozyme_trajectory.zip", processor=pooch.Unzip())
-    cache_dir = pathlib.Path(
-        pooch.os_cache("openfe") / "t4_lysozyme_trajectory.zip.unzip/t4_lysozyme_trajectory"
-    )
-    universe = mda.Universe(
-        str(cache_dir / "t4_toluene_complex.pdb"),
-        str(cache_dir / "t4_toluene_complex.xtc"),
-    )
-    return universe
-
-
@pytest.fixture
 def beta_barrel_universe():
    with resources.as_file(resources.files("openfe.tests.data")) as d:
--- a/src/openfe/tests/protocols/restraints/test_omm_restraints.py
+++ b/src/openfe/tests/protocols/restraints/test_omm_restraints.py
@@ -2,10 +2,8 @@
 # For details, see https://github.com/OpenFreeEnergy/openfe

 import os
-import pathlib

 import openmm
-import pooch
 import pytest
 from gufe import SmallMoleculeComponent
 from openff.units import unit
@@ -28,7 +26,7 @@ from openfe.protocols.restraint_utils.settings import (
    FlatBottomRestraintSettings,
 )

-from ...conftest import HAS_INTERNET
+from ...conftest import HAS_INTERNET, POOCH_CACHE


 def test_parameter_state_default():
@@ -98,34 +96,18 @@ def test_verify_geometry():
        restraint._verify_geometry(geometry)


-POOCH_CACHE = pooch.os_cache("openfe")
-zenodo_restraint_data = pooch.create(
-    path=POOCH_CACHE,
-    base_url="doi:10.5281/zenodo.15212342",
-    registry={
-        "industry_benchmark_systems.zip": "sha256:2bb5eee36e29b718b96bf6e9350e0b9957a592f6c289f77330cbb6f4311a07bd"
-    },
-)
-
-
@pytest.fixture
-def tyk2_protein_ligand_system():
-    zenodo_restraint_data.fetch("industry_benchmark_systems.zip", processor=pooch.Unzip())
-    cache_dir = pathlib.Path(
-        pooch.os_cache("openfe") / "industry_benchmark_systems.zip.unzip/industry_benchmark_systems"
-    )
-    with open(str(cache_dir / "jacs_set" / "tyk2" / "protein_ligand_system.xml")) as xml:
+def tyk2_protein_ligand_system(industry_benchmark_files):
+    with open(
+        str(industry_benchmark_files / "jacs_set" / "tyk2" / "protein_ligand_system.xml")
+    ) as xml:
        return openmm.XmlSerializer.deserialize(xml.read())


@pytest.fixture
-def tyk2_rdkit_ligand():
-    zenodo_restraint_data.fetch("industry_benchmark_systems.zip", processor=pooch.Unzip())
-    cache_dir = pathlib.Path(
-        pooch.os_cache("openfe") / "industry_benchmark_systems.zip.unzip/industry_benchmark_systems"
-    )
+def tyk2_rdkit_ligand(industry_benchmark_files):
    ligand = SmallMoleculeComponent.from_sdf_file(
-        str(cache_dir / "jacs_set" / "tyk2" / "test_ligand.sdf")
+        str(industry_benchmark_files / "jacs_set" / "tyk2" / "test_ligand.sdf")
    )
    return ligand.to_rdkit()

--- a/src/openfe/tests/protocols/test_openmmutils.py
+++ b/src/openfe/tests/protocols/test_openmmutils.py
@@ -40,7 +40,7 @@ from openfe.protocols.openmm_utils.charge_generation import (
    HAS_NAGL,
    HAS_OPENEYE,
 )
-from openfe.tests.conftest import HAS_INTERNET
+from openfe.tests.conftest import HAS_INTERNET, POOCH_CACHE


@pytest.mark.parametrize(
@@ -1033,25 +1033,6 @@ class TestOFFPartialCharge:
            )


-POOCH_CACHE = pooch.os_cache("openfe")
-RFE_OUTPUT = pooch.create(
-    path=POOCH_CACHE,
-    base_url="doi:10.5281/zenodo.15375081",
-    registry={
-        "checkpoint.nc": "md5:3cfd70a4cbe463403d6ec7cca84fc31a",
-        "db.json": "md5:33c8c1a0b629a52dcc291beff59fabc6",
-        "hybrid_system.pdb": "md5:44a1e78294360037acf419b95be18fb3",
-        "simulation.nc": "md5:bc4e842b47de17704d804ae345b91599",
-        "simulation_real_time_analysis.yaml": "md5:68a7d81462c42353a91bbbe5e64fd418",
-    },
-)
-
-
-@pytest.fixture
-def simulation_nc():
-    return RFE_OUTPUT.fetch("simulation.nc")
-
-
@pytest.mark.slow
@pytest.mark.skipif(
    not os.path.exists(POOCH_CACHE) and not HAS_INTERNET,