make dask client fixture in pytest

tests/test_rf_models.py

@@ -12,12 +12,18 @@
 from botocore.config import Config
 from datacube.utils.dask import start_local_dask
 
-client = start_local_dask(n_workers=1, threads_per_worker=2)
 
 project_root = Path(__file__).parents[1]
 data_dir = f"{project_root}/tests/data/"
 
 
+@pytest.fixture(scope="module")
+def dask_client():
+    client = start_local_dask(n_workers=1, threads_per_worker=2)
+    yield client
+    client.close()
+
+
 @pytest.fixture(scope="module")
 def cultivated_model_path():
     s3_bucket = "dea-public-data-dev"
@@ -399,14 +405,15 @@ def test_cultivated_predict(
     mask_bands,
     cultivated_classes,
     input_datasets,
+    dask_client,
 ):
     cultivated = StatsCultivatedClass(
         cultivated_classes,
         cultivated_model_path,
         mask_bands,
         input_bands=cultivated_input_bands,
     )
-    client.register_plugin(cultivated.dask_worker_plugin)
+    dask_client.register_plugin(cultivated.dask_worker_plugin)
     imgs = cultivated.preprocess_predict_input(input_datasets)
     res = [cultivated.predict(img).compute() for img in imgs]
     assert (
@@ -440,14 +447,15 @@ def test_cultivated_reduce(
     mask_bands,
     cultivated_classes,
     input_datasets,
+    dask_client,
 ):
     cultivated = StatsCultivatedClass(
         cultivated_classes,
         cultivated_model_path,
         mask_bands,
         input_bands=cultivated_input_bands,
     )
-    client.register_plugin(cultivated.dask_worker_plugin)
+    dask_client.register_plugin(cultivated.dask_worker_plugin)
     res = cultivated.reduce(input_datasets)
     assert res["cultivated_class"].attrs["nodata"] == 255
     assert (

tests/test_urban_model.py

@@ -15,6 +15,13 @@
 data_dir = f"{project_root}/tests/data/"
 
 
+@pytest.fixture(scope="module")
+def dask_client():
+    client = start_local_dask(n_workers=1, threads_per_worker=2)
+    yield client
+    client.close()
+
+
 @pytest.fixture(scope="module")
 def tflite_model_path():
     s3_bucket = "dea-public-data-dev"
@@ -137,11 +144,11 @@ def test_impute_missing_values(output_classes, tflite_model_path, image_groups):
     assert (res[1][:3, :3, :] == image_groups["ga_ls8"][:3, :3, :]).all()
 
 
-def test_urban_class(output_classes, tflite_model_path, image_groups):
+def test_urban_class(output_classes, tflite_model_path, image_groups, dask_client):
     # test better than random for a prediction
     # check correctness in integration test
     stats_urban = StatsUrbanClass(output_classes, tflite_model_path)
-    client.register_plugin(stats_urban.dask_worker_plugin)
+    dask_client.register_plugin(stats_urban.dask_worker_plugin)
     input_img = stats_urban.impute_missing_values_from_group(image_groups)
     input_img[0][1, 1, :] = np.nan
     input_img[1][1, 1, :] = np.nan