add --download-only option to openfe test (#1814)

* add boilerplate

* define pooch cache at top level

* try downloading one file

* put pooch cache in utils

* use dict

* add all cli data

* keep registry data all in one place

* move downloader to openfe from openfecli

* add both registries to test

* rename for clarity

* better docstring

Author: atravitz

news/test-download-all.rst

@@ -0,0 +1,23 @@
+**Added:**
+
+* Added ``openfe test --download-only`` flag, which caches all test data stored remotely.
+
+**Changed:**
+
+* <news item>
+
+**Deprecated:**
+
+* <news item>
+
+**Removed:**
+
+* <news item>
+
+**Fixed:**
+
+* <news item>
+
+**Security:**
+
+* <news item>

src/openfe/data/__init__.py

@@ -0,0 +1,33 @@
+import pooch
+
+from ._registry import zenodo_data_registry
+
+
+def retrieve_registry_data(zenodo_registry: list[dict], path: str) -> None:
+    """Helper function for pulling all test data up-front.
+
+    Parameters
+    ----------
+    path : str
+        path to store the data - usually a pooch.os_cache instance.
+
+    """
+    downloader = pooch.DOIDownloader(progressbar=True)
+
+    def _infer_processor(fname: str):
+        if fname.endswith("tar.gz"):
+            return pooch.Untar()
+        elif fname.endswith("zip"):
+            return pooch.Unzip()
+        else:
+            return None
+
+    for d in zenodo_registry:
+        pooch.retrieve(
+            url=d["base_url"] + d["fname"],
+            known_hash=d["known_hash"],
+            fname=d["fname"],
+            processor=_infer_processor(d["fname"]),
+            downloader=downloader,
+            path=path,
+        )

src/openfe/data/_downloader.py

@@ -0,0 +1,24 @@
+import pooch
+
+POOCH_CACHE = pooch.os_cache("openfe")
+
+zenodo_rfe_simulation_nc = dict(
+    base_url="doi:10.5281/zenodo.15375081/",
+    fname="simulation.nc",
+    known_hash="md5:bc4e842b47de17704d804ae345b91599",
+)
+zenodo_t4_lysozyme_traj = dict(
+    base_url="doi:10.5281/zenodo.15212342",
+    fname="t4_lysozyme_trajectory.zip",
+    known_hash="sha256:e985d055db25b5468491e169948f641833a5fbb67a23dbb0a00b57fb7c0e59c8",
+)
+zenodo_industry_benchmark_systems = dict(
+    base_url="doi:10.5281/zenodo.15212342",
+    fname="industry_benchmark_systems.zip",
+    known_hash="sha256:2bb5eee36e29b718b96bf6e9350e0b9957a592f6c289f77330cbb6f4311a07bd",
+)
+zenodo_data_registry = [
+    zenodo_rfe_simulation_nc,
+    zenodo_t4_lysozyme_traj,
+    zenodo_industry_benchmark_systems,
+]

src/openfe/data/_registry.py

@@ -22,13 +22,12 @@
 from rdkit.Chem import AllChem
 
 import openfe
+from openfe.data._registry import POOCH_CACHE
 from openfe.protocols.openmm_rfe import RelativeHybridTopologyProtocol
 from openfe.protocols.openmm_rfe._rfe_utils.relative import HybridTopologyFactory
 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

src/openfe/tests/conftest.py

@@ -17,8 +17,12 @@
 from rdkit.Geometry import Point3D
 
 import openfe
-
-from ..conftest import POOCH_CACHE
+from openfe.data._registry import (
+    POOCH_CACHE,
+    zenodo_industry_benchmark_systems,
+    zenodo_rfe_simulation_nc,
+    zenodo_t4_lysozyme_traj,
+)
 
 
 @pytest.fixture
@@ -284,55 +288,50 @@ def septop_json() -> str:
         return f.read().decode()  # type: ignore
 
 
-zenodo_industry_benchmarks_data = pooch.create(
+pooch_industry_benchmark_systems = pooch.create(
     path=POOCH_CACHE,
-    base_url="doi:10.5281/zenodo.15212342",
+    base_url=zenodo_industry_benchmark_systems["base_url"],
     registry={
-        "industry_benchmark_systems.zip": "sha256:2bb5eee36e29b718b96bf6e9350e0b9957a592f6c289f77330cbb6f4311a07bd"
+        zenodo_industry_benchmark_systems["fname"]: zenodo_industry_benchmark_systems["known_hash"]
     },
 )
 
 
 @pytest.fixture
 def industry_benchmark_files():
-    zenodo_industry_benchmarks_data.fetch("industry_benchmark_systems.zip", processor=pooch.Unzip())
+    pooch_industry_benchmark_systems.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(
+pooch_t4_lysozyme = pooch.create(
     path=POOCH_CACHE,
-    base_url="doi:10.5281/zenodo.15212342",
-    registry={
-        "t4_lysozyme_trajectory.zip": "sha256:e985d055db25b5468491e169948f641833a5fbb67a23dbb0a00b57fb7c0e59c8"
-    },
+    base_url=zenodo_t4_lysozyme_traj["base_url"],
+    registry={zenodo_t4_lysozyme_traj["fname"]: zenodo_t4_lysozyme_traj["known_hash"]},
 )
 
 
 # session scope for downstream reuse
 @pytest.fixture(scope="session")
 def t4_lysozyme_trajectory_dir():
-    zenodo_restraint_data.fetch("t4_lysozyme_trajectory.zip", processor=pooch.Unzip())
+    pooch_t4_lysozyme.fetch("t4_lysozyme_trajectory.zip", processor=pooch.Unzip())
     cache_dir = pathlib.Path(
         POOCH_CACHE / "t4_lysozyme_trajectory.zip.unzip/t4_lysozyme_trajectory"
     )
     return cache_dir
 
 
-RFE_OUTPUT = pooch.create(
-    path=POOCH_CACHE,
-    base_url="doi:10.6084/m9.figshare.24101655",
-    registry={
-        "simulation.nc": "92361a0864d4359a75399470135f56642b72c605069a4c33dbc4be6f91f28b31",
-    },
-)
-
-
 @pytest.fixture
 def simulation_nc():
-    return RFE_OUTPUT.fetch("simulation.nc")
+    return pooch.retrieve(
+        url=zenodo_rfe_simulation_nc["base_url"] + zenodo_rfe_simulation_nc["fname"],
+        known_hash=zenodo_rfe_simulation_nc["known_hash"],
+        path=POOCH_CACHE,
+    )
 
 
 @pytest.fixture

src/openfe/tests/protocols/conftest.py

@@ -9,6 +9,7 @@
 from openff.units import unit
 from rdkit import Chem
 
+from openfe.data._registry import POOCH_CACHE
 from openfe.protocols.restraint_utils.geometry.boresch.geometry import (
     BoreschRestraintGeometry,
     find_boresch_restraint,

src/openfe/tests/protocols/restraints/test_geometry_boresch.py

@@ -10,6 +10,7 @@
 from numpy.testing import assert_equal
 from openff.units import unit
 
+from openfe.data._registry import POOCH_CACHE
 from openfe.protocols.restraint_utils.geometry.boresch.host import (
     EvaluateBoreschAtoms,
     EvaluateHostAtoms1,
@@ -24,7 +25,7 @@
     is_collinear,
 )
 
-from ...conftest import HAS_INTERNET, POOCH_CACHE
+from ...conftest import HAS_INTERNET
 
 
 @pytest.fixture

src/openfe/tests/protocols/restraints/test_geometry_boresch_host.py

@@ -11,6 +11,7 @@
 from openff.units import unit
 from rdkit import Chem
 
+from openfe.data._registry import POOCH_CACHE
 from openfe.protocols.restraint_utils.geometry.utils import (
     CentroidDistanceSort,
     FindHostAtoms,

src/openfe/tests/protocols/restraints/test_geometry_utils.py

@@ -9,6 +9,7 @@
 from openff.units import unit
 from openmmtools.states import ThermodynamicState
 
+from openfe.data._registry import POOCH_CACHE
 from openfe.protocols.restraint_utils.openmm.omm_restraints import (
     BoreschRestraint,
     BoreschRestraintGeometry,
@@ -26,7 +27,7 @@
     FlatBottomRestraintSettings,
 )
 
-from ...conftest import HAS_INTERNET, POOCH_CACHE
+from ...conftest import HAS_INTERNET
 
 
 def test_parameter_state_default():