factor out common logic

Author: TomNicholas

xarray/tests/test_async.py

@@ -1,6 +1,7 @@
 import asyncio
 import time
 from collections.abc import Iterable
+from contextlib import asynccontextmanager
 from typing import TypeVar
 
 import numpy as np
@@ -98,73 +99,77 @@ def memorystore() -> "MemoryStore":
     return memorystore
 
 
+class AsyncTimer:
+    """Context manager for timing async operations and making assertions about their execution time."""
+
+    start_time: float
+    end_time: float
+    total_time: float
+
+    @asynccontextmanager
+    async def measure(self):
+        """Measure the execution time of the async code within this context."""
+        self.start_time = time.time()
+        try:
+            yield self
+        finally:
+            self.end_time = time.time()
+            self.total_time = self.end_time - self.start_time
+
+
 @requires_zarr_v3
 @pytest.mark.asyncio
 class TestAsyncLoad:
-    N_LOADS = 10
-    LATENCY = 1.0
+    N_LOADS: int = 5
+    LATENCY: float = 1.0
+
+    def assert_time_as_expected(self, total_time: float) -> None:
+        assert total_time > self.LATENCY  # Cannot possibly be quicker than this
+        assert (
+            total_time < self.LATENCY * self.N_LOADS
+        )  # If this isn't true we're gaining nothing from async
+        assert (
+            abs(total_time - self.LATENCY) < 0.5
+        )  # Should take approximately LATENCY seconds, but allow some buffer
 
-    # TODO refactor these tests
     async def test_async_load_variable(self, memorystore):
         latencystore = LatencyStore(memorystore, latency=self.LATENCY)
         ds = xr.open_zarr(latencystore, zarr_format=3, consolidated=False, chunks=None)
 
-        # TODO add async load to Dataset and DataArray as well as to Variable
         # TODO change the syntax to `.async.load()`?
-        start_time = time.time()
-        tasks = [ds["foo"].variable.async_load() for _ in range(self.N_LOADS)]
-        results = await asyncio.gather(*tasks)
-        total_time = time.time() - start_time
+        async with AsyncTimer().measure() as timer:
+            tasks = [ds["foo"].variable.async_load() for _ in range(self.N_LOADS)]
+            results = await asyncio.gather(*tasks)
 
         for result in results:
             xrt.assert_identical(result, ds["foo"].variable.load())
 
-        assert total_time > self.LATENCY  # Cannot possibly be quicker than this
-        assert (
-            total_time < self.LATENCY * self.N_LOADS
-        )  # If this isn't true we're gaining nothing from async
-        assert (
-            abs(total_time - self.LATENCY) < 0.5
-        )  # Should take approximately LATENCY seconds, but allow some buffer
+        self.assert_time_as_expected(timer.total_time)
 
     async def test_async_load_dataarray(self, memorystore):
         latencystore = LatencyStore(memorystore, latency=self.LATENCY)
         ds = xr.open_zarr(latencystore, zarr_format=3, consolidated=False, chunks=None)
 
         # TODO change the syntax to `.async.load()`?
-        start_time = time.time()
-        tasks = [ds["foo"].async_load() for _ in range(self.N_LOADS)]
-        results = await asyncio.gather(*tasks)
-        total_time = time.time() - start_time
+        async with AsyncTimer().measure() as timer:
+            tasks = [ds["foo"].async_load() for _ in range(self.N_LOADS)]
+            results = await asyncio.gather(*tasks)
 
         for result in results:
             xrt.assert_identical(result, ds["foo"].load())
 
-        assert total_time > self.LATENCY  # Cannot possibly be quicker than this
-        assert (
-            total_time < self.LATENCY * self.N_LOADS
-        )  # If this isn't true we're gaining nothing from async
-        assert (
-            abs(total_time - self.LATENCY) < 0.5
-        )  # Should take approximately LATENCY seconds, but allow some buffer
+        self.assert_time_as_expected(timer.total_time)
 
     async def test_async_load_dataset(self, memorystore):
         latencystore = LatencyStore(memorystore, latency=self.LATENCY)
         ds = xr.open_zarr(latencystore, zarr_format=3, consolidated=False, chunks=None)
 
         # TODO change the syntax to `.async.load()`?
-        start_time = time.time()
-        tasks = [ds.async_load() for _ in range(self.N_LOADS)]
-        results = await asyncio.gather(*tasks)
-        total_time = time.time() - start_time
+        async with AsyncTimer().measure() as timer:
+            tasks = [ds.async_load() for _ in range(self.N_LOADS)]
+            results = await asyncio.gather(*tasks)
 
         for result in results:
             xrt.assert_identical(result, ds.load())
 
-        assert total_time > self.LATENCY  # Cannot possibly be quicker than this
-        assert (
-            total_time < self.LATENCY * self.N_LOADS
-        )  # If this isn't true we're gaining nothing from async
-        assert (
-            abs(total_time - self.LATENCY) < 0.5
-        )  # Should take approximately LATENCY seconds, but allow some buffer
+        self.assert_time_as_expected(timer.total_time)