From 763fb1451b9bf96c55a7572c39c88db9f92dfd23 Mon Sep 17 00:00:00 2001 From: BD Himes Date: Tue, 10 Feb 2026 16:42:11 +0200 Subject: [PATCH 1/4] Extend the records attribute for each iteration, otherwise the only re-retrievable records are the initial batch at instantiation. --- async_substrate_interface/async_substrate.py | 1 + async_substrate_interface/sync_substrate.py | 1 + 2 files changed, 2 insertions(+) diff --git a/async_substrate_interface/async_substrate.py b/async_substrate_interface/async_substrate.py index 177e0e2..2a9a49c 100644 --- a/async_substrate_interface/async_substrate.py +++ b/async_substrate_interface/async_substrate.py @@ -558,6 +558,7 @@ async def __anext__(self): self.loading_complete = True raise StopAsyncIteration + self.records.extend(next_page) # Update the buffer with the newly fetched records self._buffer = iter(next_page) return next(self._buffer) diff --git a/async_substrate_interface/sync_substrate.py b/async_substrate_interface/sync_substrate.py index 5b6db72..67517e0 100644 --- a/async_substrate_interface/sync_substrate.py +++ b/async_substrate_interface/sync_substrate.py @@ -511,6 +511,7 @@ def __next__(self): self.loading_complete = True raise StopIteration + self.records.extend(next_page) # Update the buffer with the newly fetched records self._buffer = iter(next_page) return next(self._buffer) From b740788a0c63aa947419cf245453e220dd3d8eef Mon Sep 17 00:00:00 2001 From: BD Himes Date: Tue, 10 Feb 2026 16:51:03 +0200 Subject: [PATCH 2/4] Adds new method to retrieve all records from query map results --- async_substrate_interface/async_substrate.py | 12 ++++++++++++ async_substrate_interface/sync_substrate.py | 12 ++++++++++++ 2 files changed, 24 insertions(+) diff --git a/async_substrate_interface/async_substrate.py b/async_substrate_interface/async_substrate.py index 2a9a49c..b431722 100644 --- a/async_substrate_interface/async_substrate.py +++ b/async_substrate_interface/async_substrate.py @@ -526,6 +526,18 @@ async def retrieve_next_page(self, start_key) -> list: self.last_key = result.last_key return result.records + async def retrieve_all_records(self) -> list[Any]: + """ + Retrieves all records from all subsequent pages for the AsyncQueryMapResult, + returning them as a list. + + Side effect: + The self.records list will be populated fully after running this method. + """ + async for _ in self: + pass + return self.records + def __aiter__(self): return self diff --git a/async_substrate_interface/sync_substrate.py b/async_substrate_interface/sync_substrate.py index 67517e0..1a1f433 100644 --- a/async_substrate_interface/sync_substrate.py +++ b/async_substrate_interface/sync_substrate.py @@ -482,6 +482,18 @@ def retrieve_next_page(self, start_key) -> list: self.last_key = result.last_key return result.records + def retrieve_all_records(self) -> list[Any]: + """ + Retrieves all records from all subsequent pages for the QueryMapResult, + returning them as a list. + + Side effect: + The self.records list will be populated fully after running this method. + """ + for _ in self: + pass + return self.records + def __iter__(self): return self From 430defe5bbc5f8df49001ea642a36ccc4b84aa8a Mon Sep 17 00:00:00 2001 From: BD Himes Date: Tue, 10 Feb 2026 16:58:23 +0200 Subject: [PATCH 3/4] Add unit tests --- .../asyncio_/test_substrate_interface.py | 42 +++++++++++++++++++ .../sync/test_substrate_interface.py | 42 ++++++++++++++++++- 2 files changed, 83 insertions(+), 1 deletion(-) diff --git a/tests/unit_tests/asyncio_/test_substrate_interface.py b/tests/unit_tests/asyncio_/test_substrate_interface.py index 721804b..646343d 100644 --- a/tests/unit_tests/asyncio_/test_substrate_interface.py +++ b/tests/unit_tests/asyncio_/test_substrate_interface.py @@ -7,6 +7,7 @@ from websockets.protocol import State from async_substrate_interface.async_substrate import ( + AsyncQueryMapResult, AsyncSubstrateInterface, get_async_substrate_interface, ) @@ -175,3 +176,44 @@ async def test_memory_leak(): f"Loop {i}: diff={total_diff / 1024:.2f} KiB, current={current / 1024:.2f} KiB, " f"peak={peak / 1024:.2f} KiB" ) + + +@pytest.mark.asyncio +async def test_async_query_map_result_retrieve_all_records(): + """Test that retrieve_all_records fetches all pages and returns the full record list.""" + page1 = [("key1", "val1"), ("key2", "val2")] + page2 = [("key3", "val3"), ("key4", "val4")] + page3 = [("key5", "val5")] # partial page signals loading_complete + + mock_substrate = MagicMock() + + qm = AsyncQueryMapResult( + records=list(page1), + page_size=2, + substrate=mock_substrate, + module="TestModule", + storage_function="TestStorage", + last_key="key2", + ) + + # Build mock pages: first call returns page2 (full page), second returns page3 (partial) + page2_result = AsyncQueryMapResult( + records=list(page2), + page_size=2, + substrate=mock_substrate, + last_key="key4", + ) + page3_result = AsyncQueryMapResult( + records=list(page3), + page_size=2, + substrate=mock_substrate, + last_key="key5", + ) + mock_substrate.query_map = AsyncMock(side_effect=[page2_result, page3_result]) + + result = await qm.retrieve_all_records() + + assert result == page1 + page2 + page3 + assert qm.records == page1 + page2 + page3 + assert qm.loading_complete is True + assert mock_substrate.query_map.call_count == 2 diff --git a/tests/unit_tests/sync/test_substrate_interface.py b/tests/unit_tests/sync/test_substrate_interface.py index 54a5b7d..b60e065 100644 --- a/tests/unit_tests/sync/test_substrate_interface.py +++ b/tests/unit_tests/sync/test_substrate_interface.py @@ -1,7 +1,7 @@ import tracemalloc from unittest.mock import MagicMock -from async_substrate_interface.sync_substrate import SubstrateInterface +from async_substrate_interface.sync_substrate import SubstrateInterface, QueryMapResult from async_substrate_interface.types import ScaleObj from tests.helpers.settings import ARCHIVE_ENTRYPOINT, LATENT_LITE_ENTRYPOINT @@ -122,3 +122,43 @@ def test_memory_leak(): f"Loop {i}: diff={total_diff / 1024:.2f} KiB, current={current / 1024:.2f} KiB, " f"peak={peak / 1024:.2f} KiB" ) + + +def test_async_query_map_result_retrieve_all_records(): + """Test that retrieve_all_records fetches all pages and returns the full record list.""" + page1 = [("key1", "val1"), ("key2", "val2")] + page2 = [("key3", "val3"), ("key4", "val4")] + page3 = [("key5", "val5")] # partial page signals loading_complete + + mock_substrate = MagicMock() + + qm = QueryMapResult( + records=list(page1), + page_size=2, + substrate=mock_substrate, + module="TestModule", + storage_function="TestStorage", + last_key="key2", + ) + + # Build mock pages: first call returns page2 (full page), second returns page3 (partial) + page2_result = QueryMapResult( + records=list(page2), + page_size=2, + substrate=mock_substrate, + last_key="key4", + ) + page3_result = QueryMapResult( + records=list(page3), + page_size=2, + substrate=mock_substrate, + last_key="key5", + ) + mock_substrate.query_map = MagicMock(side_effect=[page2_result, page3_result]) + + result = qm.retrieve_all_records() + + assert result == page1 + page2 + page3 + assert qm.records == page1 + page2 + page3 + assert qm.loading_complete is True + assert mock_substrate.query_map.call_count == 2 From 8f64dbe40d2bb305330a81788c15b1b7c31ec5b6 Mon Sep 17 00:00:00 2001 From: BD Himes Date: Thu, 12 Feb 2026 20:04:58 +0200 Subject: [PATCH 4/4] Ruff --- tests/unit_tests/asyncio_/test_substrate_interface.py | 2 ++ tests/unit_tests/sync/test_substrate_interface.py | 2 ++ 2 files changed, 4 insertions(+) diff --git a/tests/unit_tests/asyncio_/test_substrate_interface.py b/tests/unit_tests/asyncio_/test_substrate_interface.py index c97773e..632af81 100644 --- a/tests/unit_tests/asyncio_/test_substrate_interface.py +++ b/tests/unit_tests/asyncio_/test_substrate_interface.py @@ -217,6 +217,8 @@ async def test_async_query_map_result_retrieve_all_records(): assert qm.records == page1 + page2 + page3 assert qm.loading_complete is True assert mock_substrate.query_map.call_count == 2 + + class TestGetBlockHash: @pytest.fixture def substrate(self): diff --git a/tests/unit_tests/sync/test_substrate_interface.py b/tests/unit_tests/sync/test_substrate_interface.py index 7703168..94a43a0 100644 --- a/tests/unit_tests/sync/test_substrate_interface.py +++ b/tests/unit_tests/sync/test_substrate_interface.py @@ -162,6 +162,8 @@ def test_async_query_map_result_retrieve_all_records(): assert qm.records == page1 + page2 + page3 assert qm.loading_complete is True assert mock_substrate.query_map.call_count == 2 + + class TestGetBlockHash: def _make_substrate(self): s = SubstrateInterface("ws://localhost", _mock=True)