From 1b03d4b4f59f806d180aca525fca2f63de99e4f4 Mon Sep 17 00:00:00 2001 From: Gabe Joseph Date: Mon, 31 Jan 2022 17:14:12 -0700 Subject: [PATCH] Test broadcast-trick logic --- stackstac/tests/test_to_dask.py | 14 +++++++++++++- 1 file changed, 13 insertions(+), 1 deletion(-) diff --git a/stackstac/tests/test_to_dask.py b/stackstac/tests/test_to_dask.py index 56ce85c..1a85bfb 100644 --- a/stackstac/tests/test_to_dask.py +++ b/stackstac/tests/test_to_dask.py @@ -6,6 +6,8 @@ import hypothesis.extra.numpy as st_np import numpy as np from rasterio import windows +import dask.core +import dask.threaded from dask.array.utils import assert_eq from stackstac.raster_spec import Bbox, RasterSpec @@ -171,7 +173,17 @@ def __setstate__(self, state): assert arr.dtype == dtype_ assert_eq(arr, results, equal_nan=True) - # TODO test broadcast-trick chunks + + # Check that entirely-empty chunks are broadcast-tricked into being tiny. + # NOTE: unfortunately, this computes the array again, which slows down tests. + # But passing a computed array into `assert_eq` would skip handy checks for chunks, meta, etc. + TestReader.opened.clear() + chunks = dask.threaded.get(arr.dask, list(dask.core.flatten(arr.__dask_keys__()))) + for chunk in chunks: + if ( + np.isnan(chunk) if np.isnan(fill_value_) else np.equal(chunk, fill_value_) + ).all(): + assert chunk.strides == (0, 0, 0, 0) @given(