From 7b8aad8c96f75d04085001f15135d6d7c7ad7cff Mon Sep 17 00:00:00 2001
From: Ian Thomas <ianthomas23@gmail.com>
Date: Tue, 11 Oct 2022 16:33:48 +0100
Subject: [PATCH] Consistent test_gpu flag (#1132)

---
 datashader/tests/benchmarks/test_canvas.py | 11 +++--------
 datashader/tests/test_dask.py              |  9 +++------
 datashader/tests/test_pandas.py            | 10 +++-------
 3 files changed, 9 insertions(+), 21 deletions(-)

diff --git a/datashader/tests/benchmarks/test_canvas.py b/datashader/tests/benchmarks/test_canvas.py
index d65bec12f..5853d3974 100644
--- a/datashader/tests/benchmarks/test_canvas.py
+++ b/datashader/tests/benchmarks/test_canvas.py
@@ -5,10 +5,7 @@
 
 import datashader as ds
 
-if "DATASHADER_TEST_GPU" in os.environ:
-    test_gpu = bool(int(os.environ["DATASHADER_TEST_GPU"]))
-else:
-    test_gpu = None
+test_gpu = bool(int(os.getenv("DATASHADER_TEST_GPU", 0)))
 
 
 @pytest.fixture
@@ -35,8 +32,7 @@ def test_points(benchmark, time_series):
     benchmark(cvs.points, time_series, 'x', 'y')
 
 
-@pytest.mark.skipif(test_gpu is None, reason="DATASHADER_TEST_GPU not in environment")
-@pytest.mark.skipif(test_gpu is False, reason="DATASHADER_TEST_GPU is set to False")
+@pytest.mark.skipif(not test_gpu, reason="DATASHADER_TEST_GPU not set")
 @pytest.mark.benchmark(group="canvas")
 def test_line_gpu(benchmark, time_series):
     from cudf import from_pandas
@@ -45,8 +41,7 @@ def test_line_gpu(benchmark, time_series):
     benchmark(cvs.line, time_series, 'x', 'y')
 
 
-@pytest.mark.skipif(test_gpu is None, reason="DATASHADER_TEST_GPU not in environment")
-@pytest.mark.skipif(test_gpu is False, reason="DATASHADER_TEST_GPU is set to False")
+@pytest.mark.skipif(not test_gpu, reason="DATASHADER_TEST_GPU not set")
 @pytest.mark.benchmark(group="canvas")
 def test_points_gpu(benchmark, time_series):
     from cudf import from_pandas
diff --git a/datashader/tests/test_dask.py b/datashader/tests/test_dask.py
index 971082f14..5abff3c70 100644
--- a/datashader/tests/test_dask.py
+++ b/datashader/tests/test_dask.py
@@ -27,10 +27,7 @@
 
 config.set(scheduler='synchronous')
 
-if "DATASHADER_TEST_GPU" in os.environ:
-    test_gpu = bool(int(os.environ["DATASHADER_TEST_GPU"]))
-else:
-    test_gpu = None
+test_gpu = bool(int(os.getenv("DATASHADER_TEST_GPU", 0)))
 
 df_pd = pd.DataFrame({'x': np.array(([0.] * 10 + [1] * 10)),
                       'y': np.array(([0.] * 5 + [1] * 5 + [0] * 5 + [1] * 5)),
@@ -62,7 +59,7 @@ def dask_DataFrame(*args, **kwargs):
     import cupy
     import dask_cudf
 
-    if test_gpu is False:
+    if not test_gpu:
         # GPU testing disabled even though cudf/cupy are available
         raise ImportError
 
@@ -109,7 +106,7 @@ def floats(n):
 
 
 def test_gpu_dependencies():
-    if test_gpu is True and cudf is None:
+    if test_gpu and cudf is None:
         pytest.fail(
             "cudf, cupy, and/or dask_cudf not available and DATASHADER_TEST_GPU=1"
         )
diff --git a/datashader/tests/test_pandas.py b/datashader/tests/test_pandas.py
index 81cc8911d..c93ea186d 100644
--- a/datashader/tests/test_pandas.py
+++ b/datashader/tests/test_pandas.py
@@ -29,10 +29,7 @@
 df_pd.at[2,'f64'] = nan
 dfs_pd = [df_pd]
 
-if "DATASHADER_TEST_GPU" in os.environ:
-    test_gpu = bool(int(os.environ["DATASHADER_TEST_GPU"]))
-else:
-    test_gpu = None
+test_gpu = bool(int(os.getenv("DATASHADER_TEST_GPU", 0)))
 
 
 try:
@@ -135,12 +132,11 @@ def values(s):
 
 
 def test_gpu_dependencies():
-    if test_gpu is True and cudf is None:
+    if test_gpu and cudf is None:
         pytest.fail("cudf and/or cupy not available and DATASHADER_TEST_GPU=1")
 
 
-@pytest.mark.skipif(test_gpu is None, reason="DATASHADER_TEST_GPU not in environment")
-@pytest.mark.skipif(test_gpu is False, reason="DATASHADER_TEST_GPU is set to False")
+@pytest.mark.skipif(not test_gpu, reason="DATASHADER_TEST_GPU not set")
 def test_cudf_concat():
     # Testing if a newer version of cuDF implements the possibility to
     # concatenate multiple columns with the same name.