Skip to content

Commit

Permalink
Consolidate some numbagg tests (#9211)
Browse files Browse the repository at this point in the history
Some minor repetiton
  • Loading branch information
max-sixty authored Jul 7, 2024
1 parent 971d71d commit 04b38a0
Showing 1 changed file with 21 additions and 30 deletions.
51 changes: 21 additions & 30 deletions xarray/tests/test_missing.py
Original file line number Diff line number Diff line change
Expand Up @@ -421,46 +421,37 @@ def test_ffill():
assert_equal(actual, expected)


def test_ffill_use_bottleneck_numbagg():
@pytest.mark.parametrize("compute_backend", [None], indirect=True)
@pytest.mark.parametrize("method", ["ffill", "bfill"])
def test_b_ffill_use_bottleneck_numbagg(method, compute_backend):
"""
bfill & ffill fail if both bottleneck and numba are disabled
"""
da = xr.DataArray(np.array([4, 5, np.nan], dtype=np.float64), dims="x")
with xr.set_options(use_bottleneck=False, use_numbagg=False):
with pytest.raises(RuntimeError):
da.ffill("x")
with pytest.raises(RuntimeError):
getattr(da, method)("x")


@requires_dask
def test_ffill_use_bottleneck_dask():
@pytest.mark.parametrize("compute_backend", [None], indirect=True)
@pytest.mark.parametrize("method", ["ffill", "bfill"])
def test_b_ffill_use_bottleneck_dask(method, compute_backend):
"""
ffill fails if both bottleneck and numba are disabled, on dask arrays
"""
da = xr.DataArray(np.array([4, 5, np.nan], dtype=np.float64), dims="x")
da = da.chunk({"x": 1})
with xr.set_options(use_bottleneck=False, use_numbagg=False):
with pytest.raises(RuntimeError):
da.ffill("x")
with pytest.raises(RuntimeError):
getattr(da, method)("x")


@requires_numbagg
@requires_dask
def test_ffill_use_numbagg_dask():
with xr.set_options(use_bottleneck=False):
da = xr.DataArray(np.array([4, 5, np.nan], dtype=np.float64), dims="x")
da = da.chunk(x=-1)
# Succeeds with a single chunk:
_ = da.ffill("x").compute()


def test_bfill_use_bottleneck():
da = xr.DataArray(np.array([4, 5, np.nan], dtype=np.float64), dims="x")
with xr.set_options(use_bottleneck=False, use_numbagg=False):
with pytest.raises(RuntimeError):
da.bfill("x")


@requires_dask
def test_bfill_use_bottleneck_dask():
@pytest.mark.parametrize("compute_backend", ["numbagg"], indirect=True)
def test_ffill_use_numbagg_dask(compute_backend):
da = xr.DataArray(np.array([4, 5, np.nan], dtype=np.float64), dims="x")
da = da.chunk({"x": 1})
with xr.set_options(use_bottleneck=False, use_numbagg=False):
with pytest.raises(RuntimeError):
da.bfill("x")
da = da.chunk(x=-1)
# Succeeds with a single chunk:
_ = da.ffill("x").compute()


@requires_bottleneck
Expand Down

0 comments on commit 04b38a0

Please sign in to comment.