From 5edb5dc0a09d2a45149cf6086a005e99258bd579 Mon Sep 17 00:00:00 2001
From: Tennessee Leeuwenburg <134973832+tennlee@users.noreply.github.com>
Date: Mon, 26 Aug 2024 23:03:56 +1000
Subject: [PATCH] Merge pull request #643 from tennlee/dask_type_checking

Dask type checking
---
 tests/categorical/test_binary.py      | 4 ++--
 tests/categorical/test_contingency.py | 2 +-
 tests/continuous/test_correlation.py  | 2 +-
 tests/continuous/test_flip_flop.py    | 2 +-
 tests/probabilty/test_brier.py        | 2 +-
 5 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/tests/categorical/test_binary.py b/tests/categorical/test_binary.py
index 1a26d007..6c6ec3bc 100644
--- a/tests/categorical/test_binary.py
+++ b/tests/categorical/test_binary.py
@@ -80,7 +80,7 @@ def test_pod_dask():
     result = probability_of_detection(fcst_mix.chunk(), obs1.chunk())
     assert isinstance(result.data, dask.array.Array)
     result = result.compute()
-    assert isinstance(result.data, np.ndarray)
+    assert isinstance(result.data, (np.ndarray, np.generic))
     xr.testing.assert_equal(result, expected_pod3)
 
 
@@ -134,7 +134,7 @@ def test_pofd_dask():
     result = probability_of_false_detection(fcst_mix.chunk(), obs0.chunk())
     assert isinstance(result.data, dask.array.Array)
     result = result.compute()
-    assert isinstance(result.data, np.ndarray)
+    assert isinstance(result.data, (np.ndarray, np.generic))
     xr.testing.assert_equal(result, expected_pofd3)
 
 
diff --git a/tests/categorical/test_contingency.py b/tests/categorical/test_contingency.py
index 9d86e385..e73a9722 100644
--- a/tests/categorical/test_contingency.py
+++ b/tests/categorical/test_contingency.py
@@ -373,7 +373,7 @@ def test_dask_if_available_categorical():
 
     # And that transformed tables are built out of computed things
     simple_counts = table.transform().get_counts()
-    assert isinstance(simple_counts["tp_count"].data, np.ndarray)
+    assert isinstance(simple_counts["tp_count"].data, (np.ndarray, np.generic))
 
     # And that transformed things get the same numbers
     assert table.false_alarm_rate() == table.transform().false_alarm_rate()
diff --git a/tests/continuous/test_correlation.py b/tests/continuous/test_correlation.py
index faa40c07..766fb2f4 100644
--- a/tests/continuous/test_correlation.py
+++ b/tests/continuous/test_correlation.py
@@ -101,5 +101,5 @@ def test_correlation_dask():
     result = pearsonr(DA3_CORR.chunk(), DA2_CORR.chunk())
     assert isinstance(result.data, dask.array.Array)
     result = result.compute()
-    assert isinstance(result.data, np.ndarray)
+    assert isinstance(result.data, (np.ndarray, np.generic))
     xr.testing.assert_allclose(result, EXP_CORR_REDUCE_ALL)
diff --git a/tests/continuous/test_flip_flop.py b/tests/continuous/test_flip_flop.py
index 4a08ae17..74f2d53a 100644
--- a/tests/continuous/test_flip_flop.py
+++ b/tests/continuous/test_flip_flop.py
@@ -414,7 +414,7 @@ def test_flip_flop_index_is_dask_compatible():
     assert isinstance(result.data, dask.array.Array)
     result = result.compute()
     xr.testing.assert_allclose(result, ntd.EXP_FFI_SUB_CASE0)
-    assert isinstance(result.data, np.ndarray)
+    assert isinstance(result.data, (np.ndarray, np.generic))
 
 
 def test_flip_flop_index_proportion_exceeding_is_dask_compatible():
diff --git a/tests/probabilty/test_brier.py b/tests/probabilty/test_brier.py
index ed3b8731..21a5e36b 100644
--- a/tests/probabilty/test_brier.py
+++ b/tests/probabilty/test_brier.py
@@ -89,7 +89,7 @@ def test_brier_score_dask():
     result = brier_score(FCST1.chunk(), OBS1.chunk())
     assert isinstance(result.data, dask.array.Array)
     result = result.compute()
-    assert isinstance(result.data, np.ndarray)
+    assert isinstance(result.data, (np.ndarray, np.generic))
     xr.testing.assert_equal(result, xr.DataArray(0.1))