refactor: minor refactoring, increase test coverage for query iterator/builder

Author: matthiaszimmermann

tests/test_async_query_iterator.py

@@ -89,3 +89,54 @@ async def test_async_iterate_entities_basic(
         # Verify all entity keys are present and unique
         result_keys = [entity.key for entity in entities]
         assert set(result_keys) == set(expected_keys)
+
+    @pytest.mark.asyncio
+    async def test_async_iterate_entities_empty(
+        self, async_arkiv_client_http: AsyncArkiv
+    ) -> None:
+        """Test iteration over empty result set."""
+        # Create some entities (to ensure the query mechanism works)
+        num_entities = 5
+        await create_test_entities(async_arkiv_client_http, num_entities)
+
+        # Query for non-existent batch_id
+        query = 'batch_id = "does-not-exist"'
+        options = QueryOptions(attributes=KEY | ATTRIBUTES, max_results_per_page=10)
+
+        # Collect results - should be empty
+        entities = []
+        async for entity in async_arkiv_client_http.arkiv.query_entities(
+            query=query, options=options
+        ):
+            entities.append(entity)
+
+        assert len(entities) == 0
+
+    @pytest.mark.asyncio
+    async def test_async_iterate_entities_max_results_across_pages(
+        self, async_arkiv_client_http: AsyncArkiv
+    ) -> None:
+        """Test max_results limits iteration across multiple pages."""
+        # Create 10 entities
+        num_entities = 10
+        batch_id, _ = await create_test_entities(async_arkiv_client_http, num_entities)
+
+        # Query with max_results=7 and page size of 3
+        # Should iterate across 3 pages (3+3+1) and stop at 7
+        query = f'batch_id = "{batch_id}"'
+        options = QueryOptions(
+            attributes=KEY | ATTRIBUTES,
+            max_results=7,
+            max_results_per_page=3,
+        )
+
+        entities = []
+        async for entity in async_arkiv_client_http.arkiv.query_entities(
+            query=query, options=options
+        ):
+            entities.append(entity)
+
+        assert len(entities) == 7
+        for entity in entities:
+            assert entity.attributes is not None
+            assert entity.attributes["batch_id"] == batch_id

tests/test_query_builder.py

@@ -758,12 +758,91 @@ def test_method_chaining(self, arkiv_client_http: Arkiv) -> None:
         builder4 = builder.at_block(12345)
         assert builder4 is builder
 
+    def test_limit(self, arkiv_client_http: Arkiv) -> None:
+        """Test .limit() restricts total results returned."""
+        batch_id, _ = create_test_entities(arkiv_client_http, 3)  # 3 names x 3 seq = 9
+
+        # Limit to 5 results
+        results = list(
+            arkiv_client_http.arkiv.select(KEY, ATTRIBUTES)
+            .where(f'batch_id = "{batch_id}"')
+            .limit(5)
+            .fetch()
+        )
+
+        assert len(results) == 5
+
+    def test_limit_with_order_by(self, arkiv_client_http: Arkiv) -> None:
+        """Test .limit() with ORDER BY returns top N sorted results."""
+        batch_id, _ = create_test_entities(arkiv_client_http, 3)  # 3 names x 3 seq = 9
+
+        # Get top 3 by sequence descending
+        results = list(
+            arkiv_client_http.arkiv.select(KEY, ATTRIBUTES)
+            .where(f'batch_id = "{batch_id}"')
+            .order_by(IntSort("sequence", DESC))
+            .limit(3)
+            .fetch()
+        )
+
+        assert len(results) == 3
+        # All should have sequence = 3 (highest)
+        for entity in results:
+            assert entity.attributes is not None
+            assert entity.attributes["sequence"] == 3
+
+    def test_limit_exceeds_total(self, arkiv_client_http: Arkiv) -> None:
+        """Test .limit() greater than total returns all results."""
+        batch_id, _ = create_test_entities(arkiv_client_http, 2)  # 2 names x 3 seq = 6
+
+        # Limit to 100 but only 6 exist
+        results = list(
+            arkiv_client_http.arkiv.select(KEY, ATTRIBUTES)
+            .where(f'batch_id = "{batch_id}"')
+            .limit(100)
+            .fetch()
+        )
+
+        assert len(results) == 6
+
+    def test_max_page_size(self, arkiv_client_http: Arkiv) -> None:
+        """Test .max_page_size() controls entities per page."""
+        batch_id, _ = create_test_entities(arkiv_client_http, 3)  # 3 names x 3 seq = 9
+
+        # Use small page size - should still get all results via pagination
+        results = list(
+            arkiv_client_http.arkiv.select(KEY, ATTRIBUTES)
+            .where(f'batch_id = "{batch_id}"')
+            .max_page_size(2)
+            .fetch()
+        )
+
+        # All 9 results should be returned despite small page size
+        assert len(results) == 9
+
+    def test_limit_and_max_page_size_combined(self, arkiv_client_http: Arkiv) -> None:
+        """Test .limit() and .max_page_size() work together correctly."""
+        batch_id, _ = create_test_entities(arkiv_client_http, 3)  # 3 names x 3 seq = 9
+
+        # Limit to 5, page size of 2 (should fetch 3 pages: 2+2+1)
+        results = list(
+            arkiv_client_http.arkiv.select(KEY, ATTRIBUTES)
+            .where(f'batch_id = "{batch_id}"')
+            .limit(5)
+            .max_page_size(2)
+            .fetch()
+        )
+
+        assert len(results) == 5
+
 
 class TestAsyncQueryBuilder:
     """Tests for async AsyncQueryBuilder fluent API."""
 
     @pytest.mark.asyncio
-    async def test_select_all_fields(self, async_arkiv_client_http: AsyncArkiv) -> None:
+    async def test_async_select_all_fields(
+        self, async_arkiv_client_http: AsyncArkiv
+    ) -> None:
         """Test async .select() with no args selects all fields."""
         # Create entities using sync client first (from fixture's underlying connection)
         # For async tests, we need to use the async client for everything

tests/test_query_iterator.py

@@ -197,3 +197,120 @@ def test_iterate_entities_exactly_page(self, arkiv_client_http: Arkiv) -> None:
         # Verify all entity keys are present and unique
         result_keys = [entity.key for entity in entities]
         assert set(result_keys) == set(expected_keys)
+
+    def test_iterate_entities_max_results_within_single_page(
+        self, arkiv_client_http: Arkiv
+    ) -> None:
+        """Test max_results limits iteration within a single page."""
+        # Create 10 entities
+        num_entities = 10
+        batch_id, _ = create_test_entities(arkiv_client_http, num_entities)
+
+        # Query with max_results=3 and large page size
+        # Should stop after 3 entities without fetching more
+        query = f'batch_id = "{batch_id}"'
+        options = QueryOptions(
+            attributes=KEY | ATTRIBUTES,
+            max_results=3,
+            max_results_per_page=100,
+        )
+
+        iterator = arkiv_client_http.arkiv.query_entities(query=query, options=options)
+        entities = list(iterator)
+
+        assert len(entities) == 3
+        for entity in entities:
+            assert entity.attributes is not None
+            assert entity.attributes["batch_id"] == batch_id
+
+    def test_iterate_entities_max_results_across_pages(
+        self, arkiv_client_http: Arkiv
+    ) -> None:
+        """Test max_results limits iteration across multiple pages."""
+        # Create 10 entities
+        num_entities = 10
+        batch_id, _ = create_test_entities(arkiv_client_http, num_entities)
+
+        # Query with max_results=7 and page size of 3
+        # Should iterate across 3 pages (3+3+1) and stop at 7
+        query = f'batch_id = "{batch_id}"'
+        options = QueryOptions(
+            attributes=KEY | ATTRIBUTES,
+            max_results=7,
+            max_results_per_page=3,
+        )
+
+        iterator = arkiv_client_http.arkiv.query_entities(query=query, options=options)
+        entities = list(iterator)
+
+        assert len(entities) == 7
+        for entity in entities:
+            assert entity.attributes is not None
+            assert entity.attributes["batch_id"] == batch_id
+
+    def test_iterate_entities_max_results_equals_total(
+        self, arkiv_client_http: Arkiv
+    ) -> None:
+        """Test max_results equal to total matching entities."""
+        # Create 5 entities
+        num_entities = 5
+        batch_id, expected_keys = create_test_entities(arkiv_client_http, num_entities)
+
+        # Query with max_results=5 (exactly matching total)
+        query = f'batch_id = "{batch_id}"'
+        options = QueryOptions(
+            attributes=KEY | ATTRIBUTES,
+            max_results=5,
+            max_results_per_page=10,
+        )
+
+        iterator = arkiv_client_http.arkiv.query_entities(query=query, options=options)
+        entities = list(iterator)
+
+        assert len(entities) == 5
+        result_keys = [entity.key for entity in entities]
+        assert set(result_keys) == set(expected_keys)
+
+    def test_iterate_entities_max_results_exceeds_total(
+        self, arkiv_client_http: Arkiv
+    ) -> None:
+        """Test max_results greater than total matching entities."""
+        # Create 5 entities
+        num_entities = 5
+        batch_id, expected_keys = create_test_entities(arkiv_client_http, num_entities)
+
+        # Query with max_results=100 (more than available)
+        query = f'batch_id = "{batch_id}"'
+        options = QueryOptions(
+            attributes=KEY | ATTRIBUTES,
+            max_results=100,
+            max_results_per_page=10,
+        )
+
+        iterator = arkiv_client_http.arkiv.query_entities(query=query, options=options)
+        entities = list(iterator)
+
+        # Should get all 5 entities (not artificially limited)
+        assert len(entities) == 5
+        result_keys = [entity.key for entity in entities]
+        assert set(result_keys) == set(expected_keys)
+
+    def test_iterate_entities_max_results_zero(self, arkiv_client_http: Arkiv) -> None:
+        """Test max_results=0 returns no entities."""
+        # Create 5 entities
+        num_entities = 5
+        batch_id, _ = create_test_entities(arkiv_client_http, num_entities)
+
+        # Query with max_results=0
+        query = f'batch_id = "{batch_id}"'
+        options = QueryOptions(
+            attributes=KEY | ATTRIBUTES,
+            max_results=0,
+            max_results_per_page=10,
+        )
+
+        iterator = arkiv_client_http.arkiv.query_entities(query=query, options=options)
+        entities = list(iterator)
+
+        # Should return no entities
+        assert len(entities) == 0