Fix county-loading in test_zonal_average (#1587)

coiled · Nov 11, 2024 · 483fe8e · 483fe8e
1 parent 5d54bba
commit 483fe8e
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 5 deletions.
diff --git a/tests/geospatial/test_zonal_average.py b/tests/geospatial/test_zonal_average.py
@@ -6,7 +6,6 @@
 
 
 def test_nwm(
-    s3,
     scale,
     client_factory,
     cluster_kwargs={
@@ -21,5 +20,5 @@ def test_nwm(
     with client_factory(
         **scale_kwargs[scale], **cluster_kwargs
     ) as client:  # noqa: F841
-        result = nwm(scale=scale, s3fs=s3)
+        result = nwm(scale=scale)
         result.compute()
diff --git a/tests/geospatial/workloads/zonal_average.py b/tests/geospatial/workloads/zonal_average.py
@@ -1,15 +1,14 @@
 from typing import Literal
 
 import flox
+import fsspec
 import numpy as np
 import rioxarray
 import xarray as xr
-from s3fs import S3FileSystem
 
 
 def nwm(
     scale: Literal["small", "medium", "large"],
-    s3fs: S3FileSystem,
 ) -> xr.DataArray:
     ds = xr.open_zarr(
         "s3://noaa-nwm-retrospective-2-1-zarr-pds/rtout.zarr", consolidated=True
@@ -26,7 +25,7 @@ def nwm(
     counties = rioxarray.open_rasterio(
         "s3://nwm-250m-us-counties/Counties_on_250m_grid.tif",
         chunks="auto",
-        opener=s3fs.open,
+        opener=fsspec.open,
     ).squeeze()
 
     # Remove any small floating point error in coordinate locations