diff --git a/tests/test_data.py b/tests/test_data.py index 952a046..3accb96 100644 --- a/tests/test_data.py +++ b/tests/test_data.py @@ -54,6 +54,7 @@ def create_zarr( k=0, ensemble=None, grids=None, + missing=False, ): root = zarr.group() @@ -104,6 +105,18 @@ def create_zarr( root.attrs["data_request"] = {"grid": 1, "area": "g", "param_level": {}} + if missing: + missing_dates = [] + i = 2 + p, q = 5, 7 + while i < len(dates): + missing_dates.append(dates[i]) + i = p * q + 1 + p, q = q, i + root.attrs["missing_dates"] = [ + d.astype(object).isoformat() for d in missing_dates + ] + root.create_dataset( "mean", data=np.mean(data, axis=0), @@ -157,6 +170,7 @@ def zarr_from_str(name, mode): k=int(args["k"]), ensemble=int(args["ensemble"]) if args["ensemble"] is not None else None, grids=int(args["grids"]) if args["grids"] is not None else None, + missing=args["test"] == "missing", ) @@ -199,10 +213,25 @@ def _(x): return np.array(args) +def missing(x): + if isinstance(x, tuple): + return (missing(a) for a in x) + if isinstance(x, list): + return [missing(a) for a in x] + if isinstance(x, dict): + return {k: missing(v) for k, v in x.items()} + if isinstance(x, str) and x.startswith("test-"): + return x.replace("test-", "missing-") + return x + + class DatasetTester: def __init__(self, *args, **kwargs): self.ds = open_dataset(*args, **kwargs) + args, kwargs = missing((args, kwargs)) + self.ds_missing = open_dataset(*args, **kwargs) + def run( self, *, @@ -256,6 +285,11 @@ def run( self.metadata(self.ds) assert self.ds.valid_ranges(2) == [(0, len(self.ds))] + self.missing(self.ds_missing) + + def missing(self, ds): + assert ds.missing, ds + def metadata(self, ds): metadata = ds.metadata() assert isinstance(metadata, dict)