Update

coiled · Sep 12, 2024 · 1ed67cd · 1ed67cd
1 parent d0c30d3
commit 1ed67cd
Showing 1 changed file with 4 additions and 11 deletions.
diff --git a/tests/geospatial/test_zonal_average.py b/tests/geospatial/test_zonal_average.py
@@ -4,15 +4,13 @@
 
 import dask
 import flox.xarray
-import fsspec
 import numpy as np
 import rioxarray
 import xarray as xr
 
 
 def test_nwm(
     s3_url,
-    s3_factory,
     scale,
     client_factory,
     cluster_kwargs={
@@ -22,18 +20,14 @@ def test_nwm(
     },
     scale_n_workers={
         "small": 10,
-        "large": 100,
+        "large": 200,
     },
 ):
     with client_factory(
         n_workers=scale_n_workers[scale], **cluster_kwargs
     ) as client:  # noqa: F841
         ds = xr.open_zarr(
-            fsspec.get_mapper(
-                "s3://noaa-nwm-retrospective-2-1-zarr-pds/rtout.zarr", anon=True
-            ),
-            consolidated=True,
-            chunks={"time": 896, "x": 350, "y": 350},
+            "s3://noaa-nwm-retrospective-2-1-zarr-pds/rtout.zarr", consolidated=True
         )
 
         if scale == "small":
@@ -44,10 +38,9 @@ def test_nwm(
             time_range = slice("1979-02-01", "2020-12-31")
         subset = ds.zwattablrt.sel(time=time_range)
 
-        fs = s3_factory(requester_pays=True)
         with dask.annotate(retries=3):
             counties = rioxarray.open_rasterio(
-                fs.open("s3://nwm-250m-us-counties/Counties_on_250m_grid.tif"),
+                "s3://nwm-250m-us-counties/Counties_on_250m_grid.tif",
                 chunks="auto",
             ).squeeze()
 
@@ -70,4 +63,4 @@ def test_nwm(
         county_mean.load()
         yearly_mean = county_mean.mean("time")
         # Save dataset for further analysis
-        yearly_mean.to_netcdf(f"{s3_url}/mean_zwattablrt_nwm_{scale}.nc")
+        yearly_mean.to_netcdf(s3_url)