fix: prevent additional fetch when SerpResults.yield_pages is used to fetch only 1 page.

Author: adarshdigievo

serpapi/models.py

@@ -1,10 +1,8 @@
 import json
 
-from pprint import pformat
 from collections import UserDict
 
 from .textui import prettify_json
-from .exceptions import HTTPError
 
 
 class SerpResults(UserDict):
@@ -71,6 +69,8 @@ def yield_pages(self, max_pages=1_000):
         while current_page and current_page_count < max_pages:
             yield current_page
             current_page_count += 1
+            if current_page_count >= max_pages:
+                break
             if current_page.next_page_url:
                 current_page = current_page.next_page()
             else:

tests/test_pagination.py

@@ -0,0 +1,39 @@
+import json
+from unittest.mock import Mock
+
+import pytest
+import requests
+
+import serpapi
+
+
+@pytest.mark.parametrize(
+    ("max_pages", "available_pages", "expected_pages"),
+    [(1, 3, 1), (2, 3, 2), (3, 3, 3), (5, 3, 3), (5, 1, 1), (0, 3, 0)],
+)
+def test_yield_pages_does_not_request_unused_pages(
+    max_pages, available_pages, expected_pages
+):
+    responses = []
+    for page_number in range(1, available_pages + 1):
+        data = {"search_information": {"page_number": page_number}}
+        if page_number < available_pages:
+            data["serpapi_pagination"] = {
+                "next": f"https://serpapi.com/search?engine=google&q=Coffee&start={page_number * 10}"
+            }
+        response = requests.Response()
+        response.status_code = 200
+        response.headers["Content-Type"] = "application/json"
+        response._content = json.dumps(data).encode("utf-8")
+        responses.append(response)
+
+    client = serpapi.Client(api_key="test-api-key")
+    client.session.request = Mock(side_effect=responses)
+    results = client.search(engine="google", q="Coffee")
+
+    pages = list(results.yield_pages(max_pages=max_pages))
+
+    assert [page["search_information"]["page_number"] for page in pages] == list(
+        range(1, expected_pages + 1)
+    )
+    assert client.session.request.call_count == max(1, expected_pages)