diff --git a/brainles_preprocessing/utils/zenodo.py b/brainles_preprocessing/utils/zenodo.py index b43f9d7..fe7183f 100644 --- a/brainles_preprocessing/utils/zenodo.py +++ b/brainles_preprocessing/utils/zenodo.py @@ -4,7 +4,7 @@ import zipfile from io import BytesIO from pathlib import Path -from typing import Dict, List, Tuple +from typing import ClassVar, Dict, List, Tuple import requests from loguru import logger @@ -56,6 +56,9 @@ class ZenodoException(Exception): class ZenodoRecord: BASE_URL = "https://zenodo.org/api/records" + # Process-level cache: maps (record_id, target_dir) -> resolved Path + _cache: ClassVar[Dict[Tuple[str, str], Path]] = {} + def __init__( self, record_id: str, @@ -66,8 +69,34 @@ def __init__( self.target_dir = target_dir self.label = label + @classmethod + def clear_cache(cls) -> None: + """Clear the process-level fetch cache. Primarily intended for testing.""" + cls._cache.clear() + + def _cache_key(self) -> Tuple[str, str]: + return (self.record_id, str(self.target_dir)) + def fetch(self) -> Path: - """Fetch the latest version of the record from Zenodo or from local storage.""" + """Fetch the latest version of the record from Zenodo or from local storage. + + Results are cached for the lifetime of the process so that repeated calls + (e.g. when processing many subjects in a loop) do not trigger redundant + Zenodo API requests. + """ + key = self._cache_key() + + cached = ZenodoRecord._cache.get(key) + if cached is not None: + logger.debug(f"Using cached {self.label} path: {cached}") + return cached + + result = self._fetch_uncached() + ZenodoRecord._cache[key] = result + return result + + def _fetch_uncached(self) -> Path: + """Perform the actual Zenodo check / download without consulting the cache.""" zenodo_response = self._get_metadata_and_archive_url() pattern = self._glob_pattern() diff --git a/tests/test_zenodo.py b/tests/test_zenodo.py index 7815b22..724cfe6 100644 --- a/tests/test_zenodo.py +++ b/tests/test_zenodo.py @@ -14,6 +14,14 @@ # ---- Fixtures ---- +@pytest.fixture(autouse=True) +def clear_zenodo_cache(): + """Ensure the process-level ZenodoRecord cache is clean before every test.""" + ZenodoRecord.clear_cache() + yield + ZenodoRecord.clear_cache() + + @pytest.fixture def dummy_metadata(): return { @@ -163,6 +171,76 @@ def test_fetch_replaces_old_version( mock_download.assert_called_once() +# ---- Tests for process-level caching ---- + + +@patch.object(ZenodoRecord, "_fetch_uncached") +def test_fetch_caches_result(mock_fetch_uncached, tmp_path): + """Second call to fetch() returns cached result without hitting _fetch_uncached.""" + expected_path = tmp_path / "123_v1.2.3" + mock_fetch_uncached.return_value = expected_path + + record = ZenodoRecord("123", tmp_path, "test") + + result1 = record.fetch() + result2 = record.fetch() + + assert result1 == expected_path + assert result2 == expected_path + # _fetch_uncached should only be called once despite two fetch() calls + mock_fetch_uncached.assert_called_once() + + +@patch.object(ZenodoRecord, "_fetch_uncached") +def test_fetch_cache_shared_across_instances(mock_fetch_uncached, tmp_path): + """Two separate ZenodoRecord instances with the same record_id share the cache.""" + expected_path = tmp_path / "123_v1.2.3" + mock_fetch_uncached.return_value = expected_path + + record1 = ZenodoRecord("123", tmp_path, "test") + record2 = ZenodoRecord("123", tmp_path, "other_label") + + result1 = record1.fetch() + result2 = record2.fetch() + + assert result1 == expected_path + assert result2 == expected_path + mock_fetch_uncached.assert_called_once() + + +@patch.object(ZenodoRecord, "_fetch_uncached") +def test_fetch_cache_different_records_independent(mock_fetch_uncached, tmp_path): + """Different record_ids each trigger their own fetch.""" + path_a = tmp_path / "aaa_v1.0.0" + path_b = tmp_path / "bbb_v2.0.0" + mock_fetch_uncached.side_effect = [path_a, path_b] + + record_a = ZenodoRecord("aaa", tmp_path, "a") + record_b = ZenodoRecord("bbb", tmp_path, "b") + + result_a = record_a.fetch() + result_b = record_b.fetch() + + assert result_a == path_a + assert result_b == path_b + assert mock_fetch_uncached.call_count == 2 + + +def test_clear_cache_resets_state(tmp_path): + """clear_cache() forces a fresh Zenodo check on the next fetch() call.""" + with patch.object(ZenodoRecord, "_fetch_uncached") as mock_fetch_uncached: + expected_path = tmp_path / "123_v1.2.3" + mock_fetch_uncached.return_value = expected_path + + record = ZenodoRecord("123", tmp_path, "test") + record.fetch() + assert mock_fetch_uncached.call_count == 1 + + ZenodoRecord.clear_cache() + record.fetch() + assert mock_fetch_uncached.call_count == 2 + + # ---- fetch_atlases and fetch_synthstrip ----