From e44a27a9348af8a7f643a98cfda862efd8bbce87 Mon Sep 17 00:00:00 2001
From: Tres Seaver <tseaver@palladion.com>
Date: Mon, 3 May 2021 17:45:01 -0400
Subject: [PATCH] tests: add coverage for skip < offset case

Closes #158
---
 google/cloud/datastore/query.py |  5 +++
 tests/unit/test_query.py        | 60 +++++++++++++++++++++++++++++++--
 2 files changed, 63 insertions(+), 2 deletions(-)

diff --git a/google/cloud/datastore/query.py b/google/cloud/datastore/query.py
index 2f455b6f..b4b24ca7 100644
--- a/google/cloud/datastore/query.py
+++ b/google/cloud/datastore/query.py
@@ -593,8 +593,12 @@ def _next_page(self):
             # skipped all of the results yet. Don't return any results.
             # Instead, rerun query, adjusting offsets. Datastore doesn't process
             # more than 1000 skipped results in a query.
+            old_query_pb = query_pb
+            query_pb = query_pb2.Query()
+            query_pb._pb.CopyFrom(old_query_pb._pb)  # copy for testability
             query_pb.start_cursor = response_pb.batch.skipped_cursor
             query_pb.offset -= response_pb.batch.skipped_results
+
             response_pb = self.client._datastore_api.run_query(
                 request={
                     "project_id": self._query.project,
@@ -604,6 +608,7 @@ def _next_page(self):
                 },
                 **kwargs,
             )
+
         entity_pbs = self._process_query_results(response_pb)
         return page_iterator.Page(self, entity_pbs, self.item_to_value)
 
diff --git a/tests/unit/test_query.py b/tests/unit/test_query.py
index 97c1db88..89bc7e2c 100644
--- a/tests/unit/test_query.py
+++ b/tests/unit/test_query.py
@@ -609,6 +609,58 @@ def test__next_page_no_more(self):
         self.assertIsNone(page)
         ds_api.run_query.assert_not_called()
 
+    def test__next_page_w_skipped_lt_offset(self):
+        from google.api_core import page_iterator
+        from google.cloud.datastore_v1.types import datastore as datastore_pb2
+        from google.cloud.datastore_v1.types import entity as entity_pb2
+        from google.cloud.datastore_v1.types import query as query_pb2
+        from google.cloud.datastore.query import Query
+
+        project = "prujekt"
+        skipped_1 = 100
+        skipped_cursor_1 = b"DEADBEEF"
+        skipped_2 = 50
+        skipped_cursor_2 = b"FACEDACE"
+
+        more_enum = query_pb2.QueryResultBatch.MoreResultsType.NOT_FINISHED
+
+        result_1 = _make_query_response([], b"", more_enum, skipped_1)
+        result_1.batch.skipped_cursor = skipped_cursor_1
+        result_2 = _make_query_response([], b"", more_enum, skipped_2)
+        result_2.batch.skipped_cursor = skipped_cursor_2
+
+        ds_api = _make_datastore_api(result_1, result_2)
+        client = _Client(project, datastore_api=ds_api)
+
+        query = Query(client)
+        offset = 150
+        iterator = self._make_one(query, client, offset=offset)
+
+        page = iterator._next_page()
+
+        self.assertIsInstance(page, page_iterator.Page)
+        self.assertIs(page._parent, iterator)
+
+        partition_id = entity_pb2.PartitionId(project_id=project)
+        read_options = datastore_pb2.ReadOptions()
+
+        query_1 = query_pb2.Query(offset=offset)
+        query_2 = query_pb2.Query(
+            start_cursor=skipped_cursor_1, offset=(offset - skipped_1)
+        )
+        expected_calls = [
+            mock.call(
+                request={
+                    "project_id": project,
+                    "partition_id": partition_id,
+                    "read_options": read_options,
+                    "query": query,
+                }
+            )
+            for query in [query_1, query_2]
+        ]
+        self.assertEqual(ds_api.run_query.call_args_list, expected_calls)
+
 
 class Test__item_to_entity(unittest.TestCase):
     def _call_fut(self, iterator, entity_pb):
@@ -789,6 +841,10 @@ def _make_query_response(
     )
 
 
-def _make_datastore_api(result=None):
-    run_query = mock.Mock(return_value=result, spec=[])
+def _make_datastore_api(*results):
+    if len(results) == 0:
+        run_query = mock.Mock(return_value=None, spec=[])
+    else:
+        run_query = mock.Mock(side_effect=results, spec=[])
+
     return mock.Mock(run_query=run_query, spec=["run_query"])