change xarray.tests to xarray.testing, and copy xarray.tests.raise_if…

xskillscore/tests/test_accessor_deterministic.py

@@ -1,6 +1,6 @@
 import pytest
 import xarray as xr
-from xarray.tests import assert_allclose
+from xarray.testing import assert_allclose
 
 from xskillscore.core.deterministic import (
     effective_sample_size,

xskillscore/tests/test_accessor_probabilistic.py

@@ -2,7 +2,7 @@
 import pytest
 import xarray as xr
 from scipy.stats import norm
-from xarray.tests import assert_allclose
+from xarray.testing import assert_allclose
 
 from xskillscore.core.probabilistic import (
     brier_score,

xskillscore/tests/test_deterministic.py

@@ -1,7 +1,7 @@
 import numpy as np
 import pytest
 import xarray as xr
-from xarray.tests import assert_allclose
+from xarray.testing import assert_allclose
 
 from xskillscore.core.deterministic import (
     _preprocess_dims,

xskillscore/tests/test_probabilistic.py

@@ -7,7 +7,7 @@
 from scipy.stats import norm
 from sklearn.calibration import calibration_curve
 from sklearn.metrics import roc_auc_score, roc_curve
-from xarray.tests import assert_allclose, assert_identical
+from xarray.testing import assert_allclose, assert_identical
 
 from xskillscore.core.probabilistic import (
     brier_score,

xskillscore/tests/test_skipna_functionality.py

@@ -2,7 +2,7 @@
 
 import numpy as np
 import pytest
-from xarray.tests import assert_allclose, raise_if_dask_computes
+from xarray.testing import assert_allclose
 
 from xskillscore.core.deterministic import (
     linslope,
@@ -20,6 +20,41 @@
     spearman_r_p_value,
 )
 
+from contextlib import nullcontext
+try:
+    import dask
+except ImportError:
+    has_dask = False
+else:
+    has_dask = True
+
+
+class CountingScheduler:
+    """Simple dask scheduler counting the number of computes.
+
+    Reference: https://stackoverflow.com/questions/53289286/"""
+
+    def __init__(self, max_computes=0):
+        self.total_computes = 0
+        self.max_computes = max_computes
+
+    def __call__(self, dsk, keys, **kwargs):
+        self.total_computes += 1
+        if self.total_computes > self.max_computes:
+            raise RuntimeError(
+                "Too many computes. Total: %d > max: %d."
+                % (self.total_computes, self.max_computes)
+            )
+        return dask.get(dsk, keys, **kwargs)
+
+
+def raise_if_dask_computes(max_computes=0):
+    # return a dummy context manager so that this can be used for non-dask objects
+    if not has_dask:
+        return nullcontext()
+    scheduler = CountingScheduler(max_computes)
+    return dask.config.set(scheduler=scheduler)
+
 WEIGHTED_METRICS: List[Callable] = [
     linslope,
     pearson_r,