Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
66 changes: 66 additions & 0 deletions tests/test_15_dataset_cacher.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,66 @@
from typing import Any

import numpy as np
import xarray as xr

from xarray_ecmwf.engine_ecmwf import DatasetCacher


class DummyRequestClient:
def __init__(self) -> None:
self.submit_calls = 0

def submit_and_wait_on_result(self, request: dict[str, Any]) -> Any:
self.submit_calls += 1
return {"request": request}

def get_filename(self, result: Any) -> str:
return "unused-by-request-hash-cache-key.grib"

def download(self, result: Any, target: str | None = None) -> str:
xr.Dataset({"t": ("x", np.arange(3, dtype="float32"))}).to_zarr(target, mode="w")
return target # type: ignore


def _open_zarr(path: str) -> xr.Dataset:
return xr.open_dataset(path, engine="zarr")


def test_retrieve_once_checks_the_local_cache_before_contacting_the_client(tmp_path: Any) -> None:
client = DummyRequestClient()
cacher = DatasetCacher(client, open_dataset=_open_zarr, cache_folder=str(tmp_path))
request = {"dataset": "reanalysis-era5-pressure-levels", "day": "01"}

with cacher.retrieve_once(request) as ds:
assert "t" in ds.data_vars
assert client.submit_calls == 1

with cacher.retrieve_once(request) as ds:
assert "t" in ds.data_vars
assert client.submit_calls == 1


def test_retrieve_once_still_submits_a_genuinely_different_request(tmp_path: Any) -> None:
client = DummyRequestClient()
cacher = DatasetCacher(client, open_dataset=_open_zarr, cache_folder=str(tmp_path))

with cacher.retrieve_once({"dataset": "reanalysis-era5-pressure-levels", "day": "01"}):
pass
with cacher.retrieve_once({"dataset": "reanalysis-era5-pressure-levels", "day": "02"}):
pass

assert client.submit_calls == 2


def test_retrieve_once_with_cache_file_false_still_skips_a_within_context_resubmit(tmp_path: Any) -> None:
# override_cache_file only controls whether the file is deleted after the context manager exits
# it must not affect whether an already-present file is reused.
client = DummyRequestClient()
cacher = DatasetCacher(client, open_dataset=_open_zarr, cache_folder=str(tmp_path), cache_file=False)
request = {"dataset": "reanalysis-era5-pressure-levels", "day": "01"}

with cacher.retrieve_once(request):
with cacher.retrieve_once(request):
pass

assert client.submit_calls == 1
8 changes: 4 additions & 4 deletions xarray_ecmwf/engine_ecmwf.py
Original file line number Diff line number Diff line change
Expand Up @@ -101,15 +101,15 @@ def retrieve_once(
if override_cache_file is not None:
cache_file = override_cache_file

result = self.request_client.submit_and_wait_on_result(request)
filename = self.request_client.get_filename(result)
path = os.path.join(self.cache_folder, filename)

if not os.path.isdir(self.cache_folder):
os.makedirs(self.cache_folder, exist_ok=True)

request_hash = hashlib.md5(str(request).encode("utf-8")).hexdigest()
path = os.path.join(self.cache_folder, request_hash + ".grib")

with xr.backends.locks.get_write_lock(f"{HOSTNAME}-grib"):
if not os.path.exists(path):
result = self.request_client.submit_and_wait_on_result(request)
robust_save_to_file(self.request_client.download, (result,), path)
ds = self.open_dataset(path)
LOGGER.debug("request: %r ->\n%r", request, list(ds.data_vars.values())[0])
Expand Down
Loading