Cycles 16-17: exercise max_total_rows cap + allowlist_policy edges

- New test_max_total_rows_cap.py triggers the _drain_continuations
  cap directly via mocked continuation frames: raises at exceed,
  accepts at the boundary, None disables.
- allowlist_policy now takes Iterable[str] and materializes to
  frozenset once (safe for generators, dict_keys, etc.). New tests
  cover: empty allowlist rejects all redirects; iterator input
  works across repeated calls (proves one-shot materialization).

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

Author: antoineleclair

src/dqliteclient/cluster.py

@@ -2,7 +2,7 @@
 
 import asyncio
 import random
-from collections.abc import Callable
+from collections.abc import Callable, Iterable
 
 from dqliteclient.connection import DqliteConnection, _parse_address
 from dqliteclient.exceptions import (
@@ -191,15 +191,19 @@ async def update_nodes(self, nodes: list[NodeInfo]) -> None:
         await self._node_store.set_nodes(nodes)
 
 
-def allowlist_policy(addresses: list[str] | set[str]) -> RedirectPolicy:
+def allowlist_policy(addresses: Iterable[str]) -> RedirectPolicy:
     """Build a redirect policy that accepts only the given addresses.
 
     Useful for the common case: "only allow redirects to hosts I've
     explicitly seed-listed." Addresses are matched by exact string
     equality — callers that need CIDR / DNS / wildcard matching should
     supply their own callable.
+
+    Accepts any iterable (list, set, tuple, generator, dict_keys). The
+    iterable is materialized into a frozen set once, so passing a
+    generator is safe — the returned closure doesn't re-iterate.
     """
-    allowed = set(addresses)
+    allowed = frozenset(addresses)
 
     def policy(addr: str) -> bool:
         return addr in allowed

tests/test_max_total_rows_cap.py

@@ -0,0 +1,71 @@
+"""max_total_rows cap actually fires in the continuation-drain loop.
+
+Cycle 9 wired the cap through the layers; cycle 16 adds the missing
+test that exercises the enforcement itself. Uses a mocked protocol
+response stream so we don't need a cluster to deliver millions of rows.
+"""
+
+import asyncio
+from unittest.mock import AsyncMock, MagicMock
+
+import pytest
+
+from dqliteclient.exceptions import ProtocolError
+from dqliteclient.protocol import DqliteProtocol
+
+
+def _make_rows_response(rows: list[list[object]], has_more: bool) -> MagicMock:
+    r = MagicMock(name="RowsResponse")
+    r.rows = rows
+    r.has_more = has_more
+    r.column_names = ["v"]
+    r.column_types = [1]  # INTEGER
+    return r
+
+
+class TestMaxTotalRowsEnforcement:
+    def test_exceeding_cap_raises_protocol_error(self) -> None:
+        """A continuation frame that pushes us past max_total_rows raises."""
+        reader = MagicMock()
+        writer = MagicMock()
+        p = DqliteProtocol(reader, writer, timeout=5.0, max_total_rows=5)
+
+        initial = _make_rows_response([[1], [2], [3]], has_more=True)
+        # This next frame would bring total to 6, exceeding the cap of 5.
+        over_cap = _make_rows_response([[4], [5], [6]], has_more=False)
+
+        p._read_continuation = AsyncMock(return_value=over_cap)  # type: ignore[method-assign]
+
+        async def run() -> None:
+            await p._drain_continuations(initial, deadline=999999.0)
+
+        with pytest.raises(ProtocolError, match="max_total_rows"):
+            asyncio.run(run())
+
+    def test_exactly_at_cap_does_not_raise(self) -> None:
+        """Hitting the cap exactly is fine; only exceeding raises."""
+        reader = MagicMock()
+        writer = MagicMock()
+        p = DqliteProtocol(reader, writer, timeout=5.0, max_total_rows=5)
+
+        initial = _make_rows_response([[1], [2], [3]], has_more=True)
+        at_cap = _make_rows_response([[4], [5]], has_more=False)  # total: 5
+
+        p._read_continuation = AsyncMock(return_value=at_cap)  # type: ignore[method-assign]
+
+        rows = asyncio.run(p._drain_continuations(initial, deadline=999999.0))
+        assert len(rows) == 5
+
+    def test_none_disables_cap(self) -> None:
+        """max_total_rows=None means the cap never fires."""
+        reader = MagicMock()
+        writer = MagicMock()
+        p = DqliteProtocol(reader, writer, timeout=5.0, max_total_rows=None)
+
+        initial = _make_rows_response([[i] for i in range(100)], has_more=True)
+        big = _make_rows_response([[i] for i in range(100, 10_000)], has_more=False)
+
+        p._read_continuation = AsyncMock(return_value=big)  # type: ignore[method-assign]
+
+        rows = asyncio.run(p._drain_continuations(initial, deadline=999999.0))
+        assert len(rows) == 10_000

tests/test_redirect_policy.py

@@ -52,6 +52,40 @@ def test_no_policy_means_any_redirect_accepted(self) -> None:
             result = asyncio.run(cc.find_leader())
         assert result == "anywhere.invalid:9001"
 
+    def test_empty_allowlist_rejects_all_redirects(self) -> None:
+        """Empty allowlist means every redirect fails; self-leader still
+        works because that path bypasses the policy entirely (it's not a
+        real redirect)."""
+        store = MemoryNodeStore(["10.0.0.1:9001"])
+        cc = ClusterClient(
+            store,
+            timeout=5.0,
+            redirect_policy=allowlist_policy([]),
+        )
+        with (
+            patch.object(cc, "_query_leader", new=AsyncMock(return_value="other:9001")),
+            pytest.raises(ClusterError, match="rejected"),
+        ):
+            asyncio.run(cc.find_leader())
+
+    def test_allowlist_accepts_iterator_input(self) -> None:
+        """The helper accepts any iterable; internally it materializes once
+        into a set, so a generator is safe (no iterator-exhaustion trap on
+        repeated calls)."""
+        store = MemoryNodeStore(["10.0.0.1:9001"])
+        cc = ClusterClient(
+            store,
+            timeout=5.0,
+            redirect_policy=allowlist_policy(x for x in ["10.0.0.1:9001", "10.0.0.2:9001"]),
+        )
+        # First call — would have drained the generator already.
+        with patch.object(cc, "_query_leader", new=AsyncMock(return_value="10.0.0.2:9001")):
+            assert asyncio.run(cc.find_leader()) == "10.0.0.2:9001"
+        # Second call still honors the allowlist (proves the set was
+        # materialized up-front, not re-iterated).
+        with patch.object(cc, "_query_leader", new=AsyncMock(return_value="10.0.0.2:9001")):
+            assert asyncio.run(cc.find_leader()) == "10.0.0.2:9001"
+
     def test_self_leader_bypasses_policy(self) -> None:
         """If the queried node is the leader (returns its own address),
         the redirect policy doesn't apply — the address is already in the