add --download-only option to openfe test

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/_downloader.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/_registry.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/tests/conftest.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/protocols/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/restraints/test_geometry_boresch.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_host.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_utils.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_omm_restraints.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():

src/openfe/tests/protocols/test_openmmutils.py

@@ -23,6 +23,7 @@
 from pymbar.utils import ParameterError
 
 import openfe
+from openfe.data._registry import POOCH_CACHE
 from openfe.protocols.openmm_rfe.equil_rfe_settings import (
     IntegratorSettings,
     OpenMMSolvationSettings,
@@ -40,7 +41,8 @@
     HAS_NAGL,
     HAS_OPENEYE,
 )
-from openfe.tests.conftest import HAS_INTERNET, POOCH_CACHE
+
+from ..conftest import HAS_INTERNET
 
 
 @pytest.mark.parametrize(

src/openfecli/commands/test.py

@@ -4,13 +4,23 @@
 import click
 import pytest
 
+from openfe.data import _downloader
+from openfe.data._registry import zenodo_data_registry as api_test_data_registry
 from openfecli import OFECommandPlugin
+from openfecli.data._registry import POOCH_CACHE
+from openfecli.data._registry import zenodo_data_registry as cli_test_data_registry
 from openfecli.utils import write
 
 
 @click.command("test", short_help="Run the OpenFE test suite")
 @click.option('--long', is_flag=True, default=False, help="Run additional tests (takes much longer)")  # fmt: skip
-def test(long):
+@click.option(
+    "--download-only",
+    is_flag=True,
+    default=False,
+    help="Download data to the cache if not already present (this is helpful if internet is unreliable). If all data exists in the cache, only the cache location is shown.",
+)
+def test(long, download_only):
     """
     Run the OpenFE test suite. This first checks that OpenFE is correctly
     imported, and then runs the main test suite, which should take several
@@ -22,6 +32,14 @@ def test(long):
     terminals, these show as green or yellow. Warnings are not a concern.
     However, You should not see anything that fails or errors (red).
     """
+
+    if download_only:
+        click.echo(f"Checking for test data in cache location:\n{POOCH_CACHE}")
+        _downloader.retrieve_registry_data(
+            cli_test_data_registry + api_test_data_registry, POOCH_CACHE
+        )
+        sys.exit(0)
+
     try:
         old_env = dict(os.environ)
         os.environ["OFE_SLOW_TESTS"] = str(long)

src/openfecli/data/_registry.py

@@ -0,0 +1,38 @@
+"""Registry for all remotely-stored CLI test data."""
+
+import pooch
+
+POOCH_CACHE = pooch.os_cache("openfe")
+zenodo_cmet_data = dict(
+    base_url="doi:10.5281/zenodo.15200083/",
+    fname="cmet_results.tar.gz",
+    known_hash="md5:a4ca67a907f744c696b09660dc1eb8ec",
+)
+zenodo_rbfe_serial_data = dict(
+    base_url="doi:10.5281/zenodo.15042470/",
+    fname="rbfe_results_serial_repeats.tar.gz",
+    known_hash="md5:2355ecc80e03242a4c7fcbf20cb45487",
+)
+zenodo_rbfe_parallel_data = dict(
+    base_url="doi:10.5281/zenodo.15042470/",
+    fname="rbfe_results_parallel_repeats.tar.gz",
+    known_hash="md5:ff7313e14eb6f2940c6ffd50f2192181",
+)
+zenodo_abfe_data = dict(
+    base_url="doi:10.5281/zenodo.17348229/",
+    fname="abfe_results.zip",
+    known_hash="md5:547f896e867cce61979d75b7e082f6ba",
+)
+zenodo_septop_data = dict(
+    base_url="doi:10.5281/zenodo.17435569/",
+    fname="septop_results.zip",
+    known_hash="md5:2cfa18da59a20228f5c75a1de6ec879e",
+)
+
+zenodo_data_registry = [
+    zenodo_cmet_data,
+    zenodo_rbfe_serial_data,
+    zenodo_rbfe_parallel_data,
+    zenodo_abfe_data,
+    zenodo_septop_data,
+]