From 35c60587920192722e7fb592b959028120584ffb Mon Sep 17 00:00:00 2001
From: Hendrik Makait <hendrik@coiled.io>
Date: Thu, 24 Oct 2024 10:58:39 +0200
Subject: [PATCH 1/2] Use `client.wait_for_workers` instead of
 `wait_for_workers` kwarg (#1571)

---
 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},

From 5758097a520148f51c22746776893395c8ca0177 Mon Sep 17 00:00:00 2001
From: Hendrik Makait <hendrik@coiled.io>
Date: Thu, 24 Oct 2024 10:59:20 +0200
Subject: [PATCH 2/2] Fix cleanup of object-store directories (#1570)

---
 tests/conftest.py | 15 ++++++++++++---
 1 file changed, 12 insertions(+), 3 deletions(-)

diff --git a/tests/conftest.py b/tests/conftest.py
index 95422dd6a8..97958561d0 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -669,7 +669,10 @@ def s3_url(s3, s3_scratch, test_name_uuid):
     try:
         yield url
     finally:
-        s3.rm(url, recursive=True)
+        try:
+            s3.rm(url, recursive=True)
+        except FileNotFoundError:
+            pass
 
 
 GCS_REGION = "us-central1"
@@ -698,7 +701,10 @@ def gcs_url(gcs, gcs_scratch, test_name_uuid):
     try:
         yield url
     finally:
-        gcs.rm(url, recursive=True)
+        try:
+            gcs.rm(url, recursive=True)
+        except FileNotFoundError:
+            pass
 
 
 @pytest.fixture(scope="session")
@@ -724,7 +730,10 @@ def az_url(az, az_scratch, test_name_uuid):
     try:
         yield url
     finally:
-        az.rm(url, recursive=True)
+        try:
+            az.rm(url, recursive=True)
+        except FileNotFoundError:
+            pass
 
 
 # this code was taken from pytest docs