refactor(waterdata): address PR 285 review (compact retry doc, drop _ASLEEP, inline trivial methods)

- Module docstring (L29-31): apply suggested wording — drop the
  "isn't slept off inline" / "doesn't block the call" rationale; the
  one-line escalation statement is enough.
- Drop the ``_ASLEEP = asyncio.sleep`` module-level test hook in favor
  of a direct ``await asyncio.sleep(delay)``; tests now patch
  ``chunking.asyncio.sleep`` (still scoped to the chunking module's
  asyncio binding, no extra indirection in production).
- Inline ``ChunkedCall.record(index, pair)`` into the one call site
  in ``_run.track``; the "single writer of ``_chunks``" invariant moves
  to a comment on ``self._chunks`` initialization.
- Inline ``ChunkedCall.combined()`` into ``_run``'s return; the
  ``partial_*`` bypass note moves to a comment at the return site,
  where it's more useful than buried in a removed helper's docstring.

No behavior change; 296 offline tests pass.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Author: thodson-usgs

dataretrieval/waterdata/chunking.py

@@ -26,9 +26,7 @@
 5xx, connect/read timeout) with exponential backoff + full jitter,
 honoring a server ``Retry-After`` when present. ``API_USGS_RETRIES``
 sets the cap (default 4; ``0`` disables). A ``Retry-After`` longer
-than the per-call ceiling isn't slept off inline — it escalates to
-the resumable interruption below so a multi-minute quota-window
-reset doesn't block the call.
+than the per-call ceiling escalates to a resumable interruption.
 
 Interruption: any mid-stream transient failure — 429, 5xx, or a bare
 transport error (connect/read timeout, oversize follow-up URL) — surfaces
@@ -1128,12 +1126,6 @@ def _retryable(exc: BaseException) -> tuple[bool, float | None]:
     return False, None
 
 
-# Sleep hook, indirected through a module global so tests can
-# ``monkeypatch.setattr`` it to a no-op instead of waiting for real
-# backoff. Production uses the stdlib call.
-_ASLEEP = asyncio.sleep
-
-
 def _retry_delay(exc: BaseException, attempt: int, policy: RetryPolicy) -> float | None:
     """
     Decide the backoff for a just-failed ``attempt`` (1-based), or ``None``
@@ -1206,7 +1198,7 @@ async def _retry(
             delay = _retry_delay(exc, attempt, policy)
             if delay is None:
                 raise
-            await _ASLEEP(delay)
+            await asyncio.sleep(delay)
 
 
 def _combine_chunk_frames(frames: list[pd.DataFrame]) -> pd.DataFrame:
@@ -1401,26 +1393,10 @@ def __init__(
         self.finalize = finalize
         # Completed (frame, response) pairs keyed by sub-args index; sparse
         # (gathered sub-requests complete out of order — see class docstring).
+        # ``_run``'s ``track`` closure is the only writer, so ``dict`` insertion
+        # order is completion order (relied on by :meth:`_combine_raw`).
         self._chunks: dict[int, tuple[pd.DataFrame, httpx.Response]] = {}
 
-    def record(self, index: int, pair: tuple[pd.DataFrame, httpx.Response]) -> None:
-        """
-        Record a completed sub-request's ``(frame, response)`` pair under
-        its sub-args index.
-
-        The single writer of ``self._chunks`` — used by the gather in
-        :meth:`_run` — so ``dict`` insertion order is completion order
-        (see :meth:`_combine_raw`).
-
-        Parameters
-        ----------
-        index : int
-            The sub-args index this completed pair belongs to.
-        pair : tuple of (pandas.DataFrame, httpx.Response)
-            The completed sub-request's ``(frame, response)`` pair.
-        """
-        self._chunks[index] = pair
-
     def wrap_failure(self, exc: BaseException) -> ChunkInterrupted | None:
         """
         Build the matching :class:`ChunkInterrupted` carrying this
@@ -1483,25 +1459,6 @@ def _combine_raw(self) -> tuple[pd.DataFrame, httpx.Response]:
             _combine_chunk_responses(responses, self.plan.canonical_url),
         )
 
-    def combined(self) -> tuple[pd.DataFrame, Any]:
-        """
-        Combine every recorded sub-request and apply :attr:`finalize`.
-
-        Returned by :meth:`_run` on a completed call (first run or
-        resume). The ``partial_*`` accessors deliberately do NOT route
-        through here — they return the raw :meth:`_combine_raw` snapshot
-        to stay cheap and side-effect-free.
-
-        Returns
-        -------
-        tuple of (pandas.DataFrame, finalized response)
-            The combined frame and whatever :attr:`finalize` produces —
-            a raw :class:`httpx.Response` by default, or the OGC
-            getters' type-coerced / column-arranged frame plus
-            ``BaseMetadata``.
-        """
-        return self.finalize(*self._combine_raw())
-
     @property
     def partial_frame(self) -> pd.DataFrame:
         """
@@ -1669,7 +1626,7 @@ async def track(
                 ) -> tuple[pd.DataFrame, httpx.Response]:
                     """One sub-request (with retry) + record + progress tick."""
                     result = await _retry(lambda: self.fetch(args), self.retry_policy)
-                    self.record(index, result)
+                    self._chunks[index] = result
                     if reporter is not None:
                         # Chunks finish out of order under gather, so tick the
                         # completed *count* rather than a positional index.
@@ -1710,7 +1667,10 @@ async def track(
                     interrupted, exc = first_transient
                     raise interrupted from exc
 
-        return self.combined()
+        # Apply the injected ``finalize`` to the raw combined result.
+        # ``partial_frame`` / ``partial_response`` deliberately bypass
+        # ``finalize`` to stay cheap and side-effect-free.
+        return self.finalize(*self._combine_raw())
 
 
 def multi_value_chunked(

tests/waterdata_chunking_test.py

@@ -60,8 +60,8 @@
 
 
 def _aiozero(_d):
-    """An async no-op sleep — monkeypatched over ``chunking._ASLEEP`` so
-    retry backoff doesn't actually wait in tests."""
+    """An async no-op sleep — monkeypatched over the ``chunking`` module's
+    ``asyncio.sleep`` so retry backoff doesn't actually wait in tests."""
 
     async def _noop():
         return None
@@ -1489,7 +1489,7 @@ def test_combine_chunk_responses_does_not_mutate_input_urls():
 # ---------------------------------------------------------------------------
 # Retry-with-backoff: RetryPolicy + _retryable + drivers + decorator wiring.
 # Conftest pins API_USGS_RETRIES=0, so these tests opt in explicitly and
-# patch chunking._SLEEP / chunking._ASLEEP to no-ops (no real backoff).
+# patch the chunking module's ``asyncio.sleep`` to a no-op (no real backoff).
 # ---------------------------------------------------------------------------
 
 
@@ -1601,7 +1601,7 @@ def test_retryable_skips_wrapped_midpagination_transient():
 
 
 def test_retry_transient_then_recovers(monkeypatch):
-    monkeypatch.setattr(_chunking, "_ASLEEP", _aiozero)
+    monkeypatch.setattr(_chunking.asyncio, "sleep", _aiozero)
     calls = {"n": 0}
 
     async def afn():
@@ -1616,7 +1616,7 @@ async def afn():
 
 
 def test_retry_exhausted_reraises(monkeypatch):
-    monkeypatch.setattr(_chunking, "_ASLEEP", _aiozero)
+    monkeypatch.setattr(_chunking.asyncio, "sleep", _aiozero)
     calls = {"n": 0}
 
     async def afn():
@@ -1635,7 +1635,7 @@ def _record(delay):
         slept.append(delay)
         return _aiozero(delay)
 
-    monkeypatch.setattr(_chunking, "_ASLEEP", _record)
+    monkeypatch.setattr(_chunking.asyncio, "sleep", _record)
     calls = {"n": 0}
 
     async def afn():
@@ -1654,7 +1654,7 @@ def _record(delay):
         slept.append(delay)
         return _aiozero(delay)
 
-    monkeypatch.setattr(_chunking, "_ASLEEP", _record)
+    monkeypatch.setattr(_chunking.asyncio, "sleep", _record)
     calls = {"n": 0}
 
     async def afn():
@@ -1673,7 +1673,7 @@ def test_retry_transient_then_success(monkeypatch):
     async def _noslept(_d):
         return None
 
-    monkeypatch.setattr(_chunking, "_ASLEEP", _noslept)
+    monkeypatch.setattr(_chunking.asyncio, "sleep", _noslept)
     calls = {"n": 0}
 
     async def afn():
@@ -1693,7 +1693,7 @@ def test_chunker_retries_transient_then_completes(monkeypatch):
     """A transient on one sub-request is retried transparently; the
     decorated call completes with no ChunkInterrupted."""
     monkeypatch.setenv("API_USGS_RETRIES", "3")
-    monkeypatch.setattr(_chunking, "_ASLEEP", _aiozero)
+    monkeypatch.setattr(_chunking.asyncio, "sleep", _aiozero)
     state = {"failed": False}
 
     async def fetch(args):
@@ -1713,7 +1713,7 @@ def test_chunker_exhausted_retries_still_resumable(monkeypatch):
     """When retries are exhausted the failure still surfaces as a
     resumable ChunkInterrupted — retries don't swallow the escape hatch."""
     monkeypatch.setenv("API_USGS_RETRIES", "2")
-    monkeypatch.setattr(_chunking, "_ASLEEP", _aiozero)
+    monkeypatch.setattr(_chunking.asyncio, "sleep", _aiozero)
     attempts = {"n": 0}
 
     async def fetch(args):
@@ -1737,7 +1737,7 @@ def test_async_fan_out_retries_transient_then_completes(monkeypatch):
     async def _noslept(_d):
         return None
 
-    monkeypatch.setattr(_chunking, "_ASLEEP", _noslept)
+    monkeypatch.setattr(_chunking.asyncio, "sleep", _noslept)
     state = {"failed": False}
 
     async def fetch_async(args):
@@ -1759,7 +1759,7 @@ def test_async_fan_out_surfaces_fatal_over_transient(monkeypatch):
     async def _noslept(_d):
         return None
 
-    monkeypatch.setattr(_chunking, "_ASLEEP", _noslept)
+    monkeypatch.setattr(_chunking.asyncio, "sleep", _noslept)
 
     async def fetch_async(args):
         # One chunk carries a deterministic programmer error; the rest are