From c972fefc4bc20e46e8c1289e749b89b14c9a9fc8 Mon Sep 17 00:00:00 2001
From: Hendrik Makait <hendrik@makait.com>
Date: Wed, 23 Oct 2024 18:42:52 +0200
Subject: [PATCH] Fix usage of wait_for_workers

---
 tests/geospatial/conftest.py                     | 2 ++
 tests/geospatial/test_atmospheric_circulation.py | 1 -
 tests/geospatial/test_climatology.py             | 2 --
 tests/geospatial/test_cloud_optimize.py          | 1 -
 tests/geospatial/test_rechunking.py              | 1 -
 tests/geospatial/test_satellite_filtering.py     | 1 -
 tests/geospatial/test_zonal_average.py           | 1 -
 7 files changed, 2 insertions(+), 7 deletions(-)

diff --git a/tests/geospatial/conftest.py b/tests/geospatial/conftest.py
index 7957bdbb19..38f6401d0e 100644
--- a/tests/geospatial/conftest.py
+++ b/tests/geospatial/conftest.py
@@ -41,6 +41,8 @@ def _(n_workers, env=None, **cluster_kwargs):
             if env:
                 cluster.send_private_envs(env=env)
             with cluster.get_client() as client:
+                # FIXME https://github.com/coiled/platform/issues/103
+                client.wait_for_workers(n_workers)
                 with benchmark_all(client):
                     yield client
 
diff --git a/tests/geospatial/test_atmospheric_circulation.py b/tests/geospatial/test_atmospheric_circulation.py
index 9c62d86874..d217b4659e 100644
--- a/tests/geospatial/test_atmospheric_circulation.py
+++ b/tests/geospatial/test_atmospheric_circulation.py
@@ -9,7 +9,6 @@ def test_atmospheric_circulation(
     cluster_kwargs={
         "workspace": "dask-benchmarks-gcp",
         "region": "us-central1",
-        "wait_for_workers": True,
     },
     scale_kwargs={
         "small": {"n_workers": 10},
diff --git a/tests/geospatial/test_climatology.py b/tests/geospatial/test_climatology.py
index 750ee61715..792dbe05e7 100644
--- a/tests/geospatial/test_climatology.py
+++ b/tests/geospatial/test_climatology.py
@@ -80,7 +80,6 @@ def test_rechunk_map_blocks(
     cluster_kwargs={
         "workspace": "dask-benchmarks-gcp",
         "region": "us-central1",
-        "wait_for_workers": True,
     },
     scale_kwargs={
         "small": {"n_workers": 10},
@@ -143,7 +142,6 @@ def test_highlevel_api(
     cluster_kwargs={
         "workspace": "dask-benchmarks-gcp",
         "region": "us-central1",
-        "wait_for_workers": True,
         "idle_timeout": "1h",
     },
     scale_kwargs={
diff --git a/tests/geospatial/test_cloud_optimize.py b/tests/geospatial/test_cloud_optimize.py
index afd273ade3..f3a15f5e74 100644
--- a/tests/geospatial/test_cloud_optimize.py
+++ b/tests/geospatial/test_cloud_optimize.py
@@ -9,7 +9,6 @@ def test_cloud_optimize(
     cluster_kwargs={
         "workspace": "dask-benchmarks",
         "region": "us-west-2",
-        "wait_for_workers": True,
     },
     scale_kwargs={
         "small": {"n_workers": 10},
diff --git a/tests/geospatial/test_rechunking.py b/tests/geospatial/test_rechunking.py
index 5188e0aced..5044ff6c05 100644
--- a/tests/geospatial/test_rechunking.py
+++ b/tests/geospatial/test_rechunking.py
@@ -9,7 +9,6 @@ def test_era5_rechunking(
     cluster_kwargs={
         "workspace": "dask-benchmarks-gcp",
         "region": "us-central1",
-        "wait_for_workers": True,
     },
     scale_kwargs={
         "small": {"n_workers": 10},
diff --git a/tests/geospatial/test_satellite_filtering.py b/tests/geospatial/test_satellite_filtering.py
index e2fe96cf8f..b6049e95ee 100644
--- a/tests/geospatial/test_satellite_filtering.py
+++ b/tests/geospatial/test_satellite_filtering.py
@@ -61,7 +61,6 @@ def test_satellite_filtering(
     cluster_kwargs={
         "workspace": "dask-benchmarks-azure",
         "region": "westeurope",
-        "wait_for_workers": True,
     },
     scale_kwargs={
         "small": {"n_workers": 10},
diff --git a/tests/geospatial/test_zonal_average.py b/tests/geospatial/test_zonal_average.py
index 5903f0a1e9..3eec34d2d7 100644
--- a/tests/geospatial/test_zonal_average.py
+++ b/tests/geospatial/test_zonal_average.py
@@ -15,7 +15,6 @@ def test_nwm(
     cluster_kwargs={
         "workspace": "dask-benchmarks",
         "region": "us-east-1",
-        "wait_for_workers": True,
     },
     scale_kwargs={
         "small": {"n_workers": 10},