Support for new cupy zonal tests

makepath · Feb 21, 2022 · bd2b5eb · bd2b5eb
1 parent 9829474
commit bd2b5eb
Showing 1 changed file with 7 additions and 7 deletions.
diff --git a/xrspatial/tests/test_zonal.py b/xrspatial/tests/test_zonal.py
@@ -13,7 +13,7 @@
 from xrspatial import trim
 from xrspatial import crop
 from xrspatial.zonal import regions
-from xrspatial.utils import doesnt_have_cuda
+from .general_checks import has_cuda_and_cupy
 
 
 from xrspatial.tests.general_checks import create_test_raster
@@ -173,16 +173,16 @@ def check_results(backend, df_result, expected_results_dict):
 
 @pytest.mark.parametrize("backend", ['numpy', 'dask+numpy', 'cupy'])
 def test_default_stats(backend, data_zones, data_values_2d, result_default_stats):
-    if backend == 'cupy' and doesnt_have_cuda():
-        pytest.skip("CUDA Device not Available")
+    if backend == 'cupy' and not has_cuda_and_cupy():
+        pytest.skip("Requires CUDA and CuPy")
     df_result = stats(zones=data_zones, values=data_values_2d)
     check_results(backend, df_result, result_default_stats)
 
 
 @pytest.mark.parametrize("backend", ['numpy', 'dask+numpy', 'cupy'])
 def test_zone_ids_stats(backend, data_zones, data_values_2d, result_zone_ids_stats):
-    if backend == 'cupy' and doesnt_have_cuda():
-        pytest.skip("CUDA Device not Available")
+    if backend == 'cupy' and not has_cuda_and_cupy():
+        pytest.skip("Requires CUDA and CuPy")
     zone_ids, expected_result = result_zone_ids_stats
     df_result = stats(zones=data_zones, values=data_values_2d,
                       zone_ids=zone_ids)
@@ -192,8 +192,8 @@ def test_zone_ids_stats(backend, data_zones, data_values_2d, result_zone_ids_sta
 @pytest.mark.parametrize("backend", ['numpy', 'cupy'])
 def test_custom_stats(backend, data_zones, data_values_2d, result_custom_stats):
     # ---- custom stats (NumPy and CuPy only) ----
-    if backend == 'cupy' and doesnt_have_cuda():
-        pytest.skip("CUDA Device not Available")
+    if backend == 'cupy' and not has_cuda_and_cupy():
+        pytest.skip("Requires CUDA and CuPy")
 
     custom_stats = {
         'double_sum': _double_sum,