• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1 // Copyright (C) 2019 Google LLC
2 //
3 // Licensed under the Apache License, Version 2.0 (the "License");
4 // you may not use this file except in compliance with the License.
5 // You may obtain a copy of the License at
6 //
7 //      http://www.apache.org/licenses/LICENSE-2.0
8 //
9 // Unless required by applicable law or agreed to in writing, software
10 // distributed under the License is distributed on an "AS IS" BASIS,
11 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 // See the License for the specific language governing permissions and
13 // limitations under the License.
14 
15 #include "icing/result/result-state-manager.h"
16 
17 #include "gmock/gmock.h"
18 #include "gtest/gtest.h"
19 #include "icing/file/filesystem.h"
20 #include "icing/portable/equals-proto.h"
21 #include "icing/schema/schema-store.h"
22 #include "icing/store/document-store.h"
23 #include "icing/testing/common-matchers.h"
24 #include "icing/testing/tmp-directory.h"
25 #include "icing/util/clock.h"
26 
27 namespace icing {
28 namespace lib {
29 namespace {
30 using ::icing::lib::portable_equals_proto::EqualsProto;
31 using ::testing::ElementsAre;
32 using ::testing::Eq;
33 using ::testing::Gt;
34 using ::testing::IsEmpty;
35 
CreateScoringSpec()36 ScoringSpecProto CreateScoringSpec() {
37   ScoringSpecProto scoring_spec;
38   scoring_spec.set_rank_by(ScoringSpecProto::RankingStrategy::DOCUMENT_SCORE);
39   return scoring_spec;
40 }
41 
CreateResultSpec(int num_per_page)42 ResultSpecProto CreateResultSpec(int num_per_page) {
43   ResultSpecProto result_spec;
44   result_spec.set_num_per_page(num_per_page);
45   return result_spec;
46 }
47 
CreateScoredHit(DocumentId document_id)48 ScoredDocumentHit CreateScoredHit(DocumentId document_id) {
49   return ScoredDocumentHit(document_id, kSectionIdMaskNone, /*score=*/1);
50 }
51 
52 class ResultStateManagerTest : public testing::Test {
53  protected:
SetUp()54   void SetUp() override {
55     schema_store_base_dir_ = GetTestTempDir() + "/schema_store";
56     filesystem_.CreateDirectoryRecursively(schema_store_base_dir_.c_str());
57     ICING_ASSERT_OK_AND_ASSIGN(
58         schema_store_,
59         SchemaStore::Create(&filesystem_, schema_store_base_dir_, &clock_));
60     SchemaProto schema;
61     schema.add_types()->set_schema_type("Document");
62     ICING_ASSERT_OK(schema_store_->SetSchema(std::move(schema)));
63 
64     doc_store_base_dir_ = GetTestTempDir() + "/document_store";
65     filesystem_.CreateDirectoryRecursively(doc_store_base_dir_.c_str());
66     ICING_ASSERT_OK_AND_ASSIGN(
67         DocumentStore::CreateResult result,
68         DocumentStore::Create(&filesystem_, doc_store_base_dir_, &clock_,
69                               schema_store_.get()));
70     document_store_ = std::move(result.document_store);
71   }
72 
TearDown()73   void TearDown() override {
74     filesystem_.DeleteDirectoryRecursively(doc_store_base_dir_.c_str());
75     filesystem_.DeleteDirectoryRecursively(schema_store_base_dir_.c_str());
76   }
77 
CreateResultState(const std::vector<ScoredDocumentHit> & scored_document_hits,int num_per_page)78   ResultState CreateResultState(
79       const std::vector<ScoredDocumentHit>& scored_document_hits,
80       int num_per_page) {
81     return ResultState(scored_document_hits, /*query_terms=*/{},
82                        SearchSpecProto::default_instance(), CreateScoringSpec(),
83                        CreateResultSpec(num_per_page), *document_store_);
84   }
85 
AddScoredDocument(DocumentId document_id)86   ScoredDocumentHit AddScoredDocument(DocumentId document_id) {
87     DocumentProto document;
88     document.set_namespace_("namespace");
89     document.set_uri(std::to_string(document_id));
90     document.set_schema("Document");
91     document_store_->Put(std::move(document));
92     return ScoredDocumentHit(document_id, kSectionIdMaskNone, /*score=*/1);
93   }
94 
document_store() const95   const DocumentStore& document_store() const { return *document_store_; }
96 
97  private:
98   Filesystem filesystem_;
99   std::string doc_store_base_dir_;
100   std::string schema_store_base_dir_;
101   Clock clock_;
102   std::unique_ptr<DocumentStore> document_store_;
103   std::unique_ptr<SchemaStore> schema_store_;
104 };
105 
TEST_F(ResultStateManagerTest,ShouldRankAndPaginateOnePage)106 TEST_F(ResultStateManagerTest, ShouldRankAndPaginateOnePage) {
107   ResultState original_result_state =
108       CreateResultState({AddScoredDocument(/*document_id=*/0),
109                          AddScoredDocument(/*document_id=*/1),
110                          AddScoredDocument(/*document_id=*/2)},
111                         /*num_per_page=*/10);
112 
113   ResultStateManager result_state_manager(
114       /*max_total_hits=*/std::numeric_limits<int>::max(), document_store());
115   ICING_ASSERT_OK_AND_ASSIGN(
116       PageResultState page_result_state,
117       result_state_manager.RankAndPaginate(std::move(original_result_state)));
118 
119   EXPECT_THAT(page_result_state.next_page_token, Eq(kInvalidNextPageToken));
120 
121   // Should get the original scored document hits
122   EXPECT_THAT(
123       page_result_state.scored_document_hits,
124       ElementsAre(EqualsScoredDocumentHit(CreateScoredHit(/*document_id=*/2)),
125                   EqualsScoredDocumentHit(CreateScoredHit(/*document_id=*/1)),
126                   EqualsScoredDocumentHit(CreateScoredHit(/*document_id=*/0))));
127 }
128 
TEST_F(ResultStateManagerTest,ShouldRankAndPaginateMultiplePages)129 TEST_F(ResultStateManagerTest, ShouldRankAndPaginateMultiplePages) {
130   ResultState original_result_state =
131       CreateResultState({AddScoredDocument(/*document_id=*/0),
132                          AddScoredDocument(/*document_id=*/1),
133                          AddScoredDocument(/*document_id=*/2),
134                          AddScoredDocument(/*document_id=*/3),
135                          AddScoredDocument(/*document_id=*/4)},
136                         /*num_per_page=*/2);
137 
138   ResultStateManager result_state_manager(
139       /*max_total_hits=*/std::numeric_limits<int>::max(), document_store());
140 
141   // First page, 2 results
142   ICING_ASSERT_OK_AND_ASSIGN(
143       PageResultState page_result_state1,
144       result_state_manager.RankAndPaginate(std::move(original_result_state)));
145   EXPECT_THAT(
146       page_result_state1.scored_document_hits,
147       ElementsAre(EqualsScoredDocumentHit(CreateScoredHit(/*document_id=*/4)),
148                   EqualsScoredDocumentHit(CreateScoredHit(/*document_id=*/3))));
149 
150   uint64_t next_page_token = page_result_state1.next_page_token;
151 
152   // Second page, 2 results
153   ICING_ASSERT_OK_AND_ASSIGN(PageResultState page_result_state2,
154                              result_state_manager.GetNextPage(next_page_token));
155   EXPECT_THAT(
156       page_result_state2.scored_document_hits,
157       ElementsAre(EqualsScoredDocumentHit(CreateScoredHit(/*document_id=*/2)),
158                   EqualsScoredDocumentHit(CreateScoredHit(/*document_id=*/1))));
159 
160   // Third page, 1 result
161   ICING_ASSERT_OK_AND_ASSIGN(PageResultState page_result_state3,
162                              result_state_manager.GetNextPage(next_page_token));
163   EXPECT_THAT(
164       page_result_state3.scored_document_hits,
165       ElementsAre(EqualsScoredDocumentHit(CreateScoredHit(/*document_id=*/0))));
166 
167   // No results
168   EXPECT_THAT(result_state_manager.GetNextPage(next_page_token),
169               StatusIs(libtextclassifier3::StatusCode::NOT_FOUND));
170 }
171 
TEST_F(ResultStateManagerTest,EmptyStateShouldReturnError)172 TEST_F(ResultStateManagerTest, EmptyStateShouldReturnError) {
173   ResultState empty_result_state = CreateResultState({}, /*num_per_page=*/1);
174 
175   ResultStateManager result_state_manager(
176       /*max_total_hits=*/std::numeric_limits<int>::max(), document_store());
177   EXPECT_THAT(
178       result_state_manager.RankAndPaginate(std::move(empty_result_state)),
179       StatusIs(libtextclassifier3::StatusCode::INVALID_ARGUMENT));
180 }
181 
TEST_F(ResultStateManagerTest,ShouldInvalidateOneToken)182 TEST_F(ResultStateManagerTest, ShouldInvalidateOneToken) {
183   ResultState result_state1 =
184       CreateResultState({AddScoredDocument(/*document_id=*/0),
185                          AddScoredDocument(/*document_id=*/1),
186                          AddScoredDocument(/*document_id=*/2)},
187                         /*num_per_page=*/1);
188   ResultState result_state2 =
189       CreateResultState({AddScoredDocument(/*document_id=*/3),
190                          AddScoredDocument(/*document_id=*/4),
191                          AddScoredDocument(/*document_id=*/5)},
192                         /*num_per_page=*/1);
193 
194   ResultStateManager result_state_manager(
195       /*max_total_hits=*/std::numeric_limits<int>::max(), document_store());
196   ICING_ASSERT_OK_AND_ASSIGN(
197       PageResultState page_result_state1,
198       result_state_manager.RankAndPaginate(std::move(result_state1)));
199   ICING_ASSERT_OK_AND_ASSIGN(
200       PageResultState page_result_state2,
201       result_state_manager.RankAndPaginate(std::move(result_state2)));
202 
203   result_state_manager.InvalidateResultState(
204       page_result_state1.next_page_token);
205 
206   // page_result_state1.next_page_token() shouldn't be found
207   EXPECT_THAT(
208       result_state_manager.GetNextPage(page_result_state1.next_page_token),
209       StatusIs(libtextclassifier3::StatusCode::NOT_FOUND));
210 
211   // page_result_state2.next_page_token() should still exist
212   ICING_ASSERT_OK_AND_ASSIGN(
213       page_result_state2,
214       result_state_manager.GetNextPage(page_result_state2.next_page_token));
215   EXPECT_THAT(
216       page_result_state2.scored_document_hits,
217       ElementsAre(EqualsScoredDocumentHit(CreateScoredHit(/*document_id=*/4))));
218 }
219 
TEST_F(ResultStateManagerTest,ShouldInvalidateAllTokens)220 TEST_F(ResultStateManagerTest, ShouldInvalidateAllTokens) {
221   ResultState result_state1 =
222       CreateResultState({AddScoredDocument(/*document_id=*/0),
223                          AddScoredDocument(/*document_id=*/1),
224                          AddScoredDocument(/*document_id=*/2)},
225                         /*num_per_page=*/1);
226   ResultState result_state2 =
227       CreateResultState({AddScoredDocument(/*document_id=*/3),
228                          AddScoredDocument(/*document_id=*/4),
229                          AddScoredDocument(/*document_id=*/5)},
230                         /*num_per_page=*/1);
231 
232   ResultStateManager result_state_manager(
233       /*max_total_hits=*/std::numeric_limits<int>::max(), document_store());
234   ICING_ASSERT_OK_AND_ASSIGN(
235       PageResultState page_result_state1,
236       result_state_manager.RankAndPaginate(std::move(result_state1)));
237   ICING_ASSERT_OK_AND_ASSIGN(
238       PageResultState page_result_state2,
239       result_state_manager.RankAndPaginate(std::move(result_state2)));
240 
241   result_state_manager.InvalidateAllResultStates();
242 
243   // page_result_state1.next_page_token() shouldn't be found
244   EXPECT_THAT(
245       result_state_manager.GetNextPage(page_result_state1.next_page_token),
246       StatusIs(libtextclassifier3::StatusCode::NOT_FOUND));
247 
248   // page_result_state2.next_page_token() shouldn't be found
249   EXPECT_THAT(
250       result_state_manager.GetNextPage(page_result_state2.next_page_token),
251       StatusIs(libtextclassifier3::StatusCode::NOT_FOUND));
252 }
253 
TEST_F(ResultStateManagerTest,ShouldRemoveOldestResultState)254 TEST_F(ResultStateManagerTest, ShouldRemoveOldestResultState) {
255   ResultState result_state1 =
256       CreateResultState({AddScoredDocument(/*document_id=*/0),
257                          AddScoredDocument(/*document_id=*/1)},
258                         /*num_per_page=*/1);
259   ResultState result_state2 =
260       CreateResultState({AddScoredDocument(/*document_id=*/2),
261                          AddScoredDocument(/*document_id=*/3)},
262                         /*num_per_page=*/1);
263   ResultState result_state3 =
264       CreateResultState({AddScoredDocument(/*document_id=*/4),
265                          AddScoredDocument(/*document_id=*/5)},
266                         /*num_per_page=*/1);
267 
268   ResultStateManager result_state_manager(/*max_total_hits=*/2,
269                                           document_store());
270   ICING_ASSERT_OK_AND_ASSIGN(
271       PageResultState page_result_state1,
272       result_state_manager.RankAndPaginate(std::move(result_state1)));
273   ICING_ASSERT_OK_AND_ASSIGN(
274       PageResultState page_result_state2,
275       result_state_manager.RankAndPaginate(std::move(result_state2)));
276   // Adding state 3 should cause state 1 to be removed.
277   ICING_ASSERT_OK_AND_ASSIGN(
278       PageResultState page_result_state3,
279       result_state_manager.RankAndPaginate(std::move(result_state3)));
280 
281   EXPECT_THAT(
282       result_state_manager.GetNextPage(page_result_state1.next_page_token),
283       StatusIs(libtextclassifier3::StatusCode::NOT_FOUND));
284 
285   ICING_ASSERT_OK_AND_ASSIGN(
286       page_result_state2,
287       result_state_manager.GetNextPage(page_result_state2.next_page_token));
288   EXPECT_THAT(page_result_state2.scored_document_hits,
289               ElementsAre(EqualsScoredDocumentHit(CreateScoredHit(
290                   /*document_id=*/2))));
291 
292   ICING_ASSERT_OK_AND_ASSIGN(
293       page_result_state3,
294       result_state_manager.GetNextPage(page_result_state3.next_page_token));
295   EXPECT_THAT(page_result_state3.scored_document_hits,
296               ElementsAre(EqualsScoredDocumentHit(CreateScoredHit(
297                   /*document_id=*/4))));
298 }
299 
TEST_F(ResultStateManagerTest,InvalidatedResultStateShouldDecreaseCurrentHitsCount)300 TEST_F(ResultStateManagerTest,
301        InvalidatedResultStateShouldDecreaseCurrentHitsCount) {
302   ResultState result_state1 =
303       CreateResultState({AddScoredDocument(/*document_id=*/0),
304                          AddScoredDocument(/*document_id=*/1)},
305                         /*num_per_page=*/1);
306   ResultState result_state2 =
307       CreateResultState({AddScoredDocument(/*document_id=*/2),
308                          AddScoredDocument(/*document_id=*/3)},
309                         /*num_per_page=*/1);
310   ResultState result_state3 =
311       CreateResultState({AddScoredDocument(/*document_id=*/4),
312                          AddScoredDocument(/*document_id=*/5)},
313                         /*num_per_page=*/1);
314 
315   // Add the first three states. Remember, the first page for each result state
316   // won't be cached (since it is returned immediately from RankAndPaginate).
317   // Each result state has a page size of 1 and a result set of 2 hits. So each
318   // result will take up one hit of our three hit budget.
319   ResultStateManager result_state_manager(/*max_total_hits=*/3,
320                                           document_store());
321   ICING_ASSERT_OK_AND_ASSIGN(
322       PageResultState page_result_state1,
323       result_state_manager.RankAndPaginate(std::move(result_state1)));
324   ICING_ASSERT_OK_AND_ASSIGN(
325       PageResultState page_result_state2,
326       result_state_manager.RankAndPaginate(std::move(result_state2)));
327   ICING_ASSERT_OK_AND_ASSIGN(
328       PageResultState page_result_state3,
329       result_state_manager.RankAndPaginate(std::move(result_state3)));
330 
331   // Invalidates state 2, so that the number of hits current cached should be
332   // decremented to 2.
333   result_state_manager.InvalidateResultState(
334       page_result_state2.next_page_token);
335 
336   // If invalidating state 2 correctly decremented the current hit count to 2,
337   // then adding state 4 should still be within our budget and no other result
338   // states should be evicted.
339   ResultState result_state4 =
340       CreateResultState({AddScoredDocument(/*document_id=*/6),
341                          AddScoredDocument(/*document_id=*/7)},
342                         /*num_per_page=*/1);
343   ICING_ASSERT_OK_AND_ASSIGN(
344       PageResultState page_result_state4,
345       result_state_manager.RankAndPaginate(std::move(result_state4)));
346 
347   ICING_ASSERT_OK_AND_ASSIGN(
348       page_result_state1,
349       result_state_manager.GetNextPage(page_result_state1.next_page_token));
350   EXPECT_THAT(page_result_state1.scored_document_hits,
351               ElementsAre(EqualsScoredDocumentHit(CreateScoredHit(
352                   /*document_id=*/0))));
353 
354   EXPECT_THAT(
355       result_state_manager.GetNextPage(page_result_state2.next_page_token),
356       StatusIs(libtextclassifier3::StatusCode::NOT_FOUND));
357 
358   ICING_ASSERT_OK_AND_ASSIGN(
359       page_result_state3,
360       result_state_manager.GetNextPage(page_result_state3.next_page_token));
361   EXPECT_THAT(page_result_state3.scored_document_hits,
362               ElementsAre(EqualsScoredDocumentHit(CreateScoredHit(
363                   /*document_id=*/4))));
364 
365   ICING_ASSERT_OK_AND_ASSIGN(
366       page_result_state4,
367       result_state_manager.GetNextPage(page_result_state4.next_page_token));
368   EXPECT_THAT(page_result_state4.scored_document_hits,
369               ElementsAre(EqualsScoredDocumentHit(CreateScoredHit(
370                   /*document_id=*/6))));
371 }
372 
TEST_F(ResultStateManagerTest,InvalidatedAllResultStatesShouldResetCurrentHitCount)373 TEST_F(ResultStateManagerTest,
374        InvalidatedAllResultStatesShouldResetCurrentHitCount) {
375   ResultState result_state1 =
376       CreateResultState({AddScoredDocument(/*document_id=*/0),
377                          AddScoredDocument(/*document_id=*/1)},
378                         /*num_per_page=*/1);
379   ResultState result_state2 =
380       CreateResultState({AddScoredDocument(/*document_id=*/2),
381                          AddScoredDocument(/*document_id=*/3)},
382                         /*num_per_page=*/1);
383   ResultState result_state3 =
384       CreateResultState({AddScoredDocument(/*document_id=*/4),
385                          AddScoredDocument(/*document_id=*/5)},
386                         /*num_per_page=*/1);
387 
388   // Add the first three states. Remember, the first page for each result state
389   // won't be cached (since it is returned immediately from RankAndPaginate).
390   // Each result state has a page size of 1 and a result set of 2 hits. So each
391   // result will take up one hit of our three hit budget.
392   ResultStateManager result_state_manager(/*max_total_hits=*/3,
393                                           document_store());
394   ICING_ASSERT_OK_AND_ASSIGN(
395       PageResultState page_result_state1,
396       result_state_manager.RankAndPaginate(std::move(result_state1)));
397   ICING_ASSERT_OK_AND_ASSIGN(
398       PageResultState page_result_state2,
399       result_state_manager.RankAndPaginate(std::move(result_state2)));
400   ICING_ASSERT_OK_AND_ASSIGN(
401       PageResultState page_result_state3,
402       result_state_manager.RankAndPaginate(std::move(result_state3)));
403 
404   // Invalidates all states so that the current hit count will be 0.
405   result_state_manager.InvalidateAllResultStates();
406 
407   // If invalidating all states correctly reset the current hit count to 0,
408   // then the entirety of state 4 should still be within our budget and no other
409   // result states should be evicted.
410   ResultState result_state4 =
411       CreateResultState({AddScoredDocument(/*document_id=*/6),
412                          AddScoredDocument(/*document_id=*/7)},
413                         /*num_per_page=*/1);
414   ResultState result_state5 =
415       CreateResultState({AddScoredDocument(/*document_id=*/8),
416                          AddScoredDocument(/*document_id=*/9)},
417                         /*num_per_page=*/1);
418   ResultState result_state6 =
419       CreateResultState({AddScoredDocument(/*document_id=*/10),
420                          AddScoredDocument(/*document_id=*/11)},
421                         /*num_per_page=*/1);
422   ICING_ASSERT_OK_AND_ASSIGN(
423       PageResultState page_result_state4,
424       result_state_manager.RankAndPaginate(std::move(result_state4)));
425   ICING_ASSERT_OK_AND_ASSIGN(
426       PageResultState page_result_state5,
427       result_state_manager.RankAndPaginate(std::move(result_state5)));
428   ICING_ASSERT_OK_AND_ASSIGN(
429       PageResultState page_result_state6,
430       result_state_manager.RankAndPaginate(std::move(result_state6)));
431 
432   EXPECT_THAT(
433       result_state_manager.GetNextPage(page_result_state1.next_page_token),
434       StatusIs(libtextclassifier3::StatusCode::NOT_FOUND));
435 
436   EXPECT_THAT(
437       result_state_manager.GetNextPage(page_result_state2.next_page_token),
438       StatusIs(libtextclassifier3::StatusCode::NOT_FOUND));
439 
440   EXPECT_THAT(
441       result_state_manager.GetNextPage(page_result_state3.next_page_token),
442       StatusIs(libtextclassifier3::StatusCode::NOT_FOUND));
443 
444   ICING_ASSERT_OK_AND_ASSIGN(
445       page_result_state4,
446       result_state_manager.GetNextPage(page_result_state4.next_page_token));
447   EXPECT_THAT(page_result_state4.scored_document_hits,
448               ElementsAre(EqualsScoredDocumentHit(CreateScoredHit(
449                   /*document_id=*/6))));
450 
451   ICING_ASSERT_OK_AND_ASSIGN(
452       page_result_state5,
453       result_state_manager.GetNextPage(page_result_state5.next_page_token));
454   EXPECT_THAT(page_result_state5.scored_document_hits,
455               ElementsAre(EqualsScoredDocumentHit(CreateScoredHit(
456                   /*document_id=*/8))));
457 
458   ICING_ASSERT_OK_AND_ASSIGN(
459       page_result_state6,
460       result_state_manager.GetNextPage(page_result_state6.next_page_token));
461   EXPECT_THAT(page_result_state6.scored_document_hits,
462               ElementsAre(EqualsScoredDocumentHit(CreateScoredHit(
463                   /*document_id=*/10))));
464 }
465 
TEST_F(ResultStateManagerTest,InvalidatedResultStateShouldDecreaseCurrentHitsCountByExactStateHitCount)466 TEST_F(
467     ResultStateManagerTest,
468     InvalidatedResultStateShouldDecreaseCurrentHitsCountByExactStateHitCount) {
469   ResultState result_state1 =
470       CreateResultState({AddScoredDocument(/*document_id=*/0),
471                          AddScoredDocument(/*document_id=*/1)},
472                         /*num_per_page=*/1);
473   ResultState result_state2 =
474       CreateResultState({AddScoredDocument(/*document_id=*/2),
475                          AddScoredDocument(/*document_id=*/3)},
476                         /*num_per_page=*/1);
477   ResultState result_state3 =
478       CreateResultState({AddScoredDocument(/*document_id=*/4),
479                          AddScoredDocument(/*document_id=*/5)},
480                         /*num_per_page=*/1);
481 
482   // Add the first three states. Remember, the first page for each result state
483   // won't be cached (since it is returned immediately from RankAndPaginate).
484   // Each result state has a page size of 1 and a result set of 2 hits. So each
485   // result will take up one hit of our three hit budget.
486   ResultStateManager result_state_manager(/*max_total_hits=*/3,
487                                           document_store());
488   ICING_ASSERT_OK_AND_ASSIGN(
489       PageResultState page_result_state1,
490       result_state_manager.RankAndPaginate(std::move(result_state1)));
491   ICING_ASSERT_OK_AND_ASSIGN(
492       PageResultState page_result_state2,
493       result_state_manager.RankAndPaginate(std::move(result_state2)));
494   ICING_ASSERT_OK_AND_ASSIGN(
495       PageResultState page_result_state3,
496       result_state_manager.RankAndPaginate(std::move(result_state3)));
497 
498   // Invalidates state 2, so that the number of hits current cached should be
499   // decremented to 2.
500   result_state_manager.InvalidateResultState(
501       page_result_state2.next_page_token);
502 
503   // If invalidating state 2 correctly decremented the current hit count to 2,
504   // then adding state 4 should still be within our budget and no other result
505   // states should be evicted.
506   ResultState result_state4 =
507       CreateResultState({AddScoredDocument(/*document_id=*/6),
508                          AddScoredDocument(/*document_id=*/7)},
509                         /*num_per_page=*/1);
510   ICING_ASSERT_OK_AND_ASSIGN(
511       PageResultState page_result_state4,
512       result_state_manager.RankAndPaginate(std::move(result_state4)));
513 
514   // If invalidating result state 2 correctly decremented the current hit count
515   // to 2 and adding state 4 correctly incremented it to 3, then adding this
516   // result state should trigger the eviction of state 1.
517   ResultState result_state5 =
518       CreateResultState({AddScoredDocument(/*document_id=*/8),
519                          AddScoredDocument(/*document_id=*/9)},
520                         /*num_per_page=*/1);
521   ICING_ASSERT_OK_AND_ASSIGN(
522       PageResultState page_result_state5,
523       result_state_manager.RankAndPaginate(std::move(result_state5)));
524 
525   EXPECT_THAT(
526       result_state_manager.GetNextPage(page_result_state1.next_page_token),
527       StatusIs(libtextclassifier3::StatusCode::NOT_FOUND));
528 
529   EXPECT_THAT(
530       result_state_manager.GetNextPage(page_result_state2.next_page_token),
531       StatusIs(libtextclassifier3::StatusCode::NOT_FOUND));
532 
533   ICING_ASSERT_OK_AND_ASSIGN(
534       page_result_state3,
535       result_state_manager.GetNextPage(page_result_state3.next_page_token));
536   EXPECT_THAT(page_result_state3.scored_document_hits,
537               ElementsAre(EqualsScoredDocumentHit(CreateScoredHit(
538                   /*document_id=*/4))));
539 
540   ICING_ASSERT_OK_AND_ASSIGN(
541       page_result_state4,
542       result_state_manager.GetNextPage(page_result_state4.next_page_token));
543   EXPECT_THAT(page_result_state4.scored_document_hits,
544               ElementsAre(EqualsScoredDocumentHit(CreateScoredHit(
545                   /*document_id=*/6))));
546 
547   ICING_ASSERT_OK_AND_ASSIGN(
548       page_result_state5,
549       result_state_manager.GetNextPage(page_result_state5.next_page_token));
550   EXPECT_THAT(page_result_state5.scored_document_hits,
551               ElementsAre(EqualsScoredDocumentHit(CreateScoredHit(
552                   /*document_id=*/8))));
553 }
554 
TEST_F(ResultStateManagerTest,GetNextPageShouldDecreaseCurrentHitsCount)555 TEST_F(ResultStateManagerTest, GetNextPageShouldDecreaseCurrentHitsCount) {
556   ResultState result_state1 =
557       CreateResultState({AddScoredDocument(/*document_id=*/0),
558                          AddScoredDocument(/*document_id=*/1)},
559                         /*num_per_page=*/1);
560   ResultState result_state2 =
561       CreateResultState({AddScoredDocument(/*document_id=*/2),
562                          AddScoredDocument(/*document_id=*/3)},
563                         /*num_per_page=*/1);
564   ResultState result_state3 =
565       CreateResultState({AddScoredDocument(/*document_id=*/4),
566                          AddScoredDocument(/*document_id=*/5)},
567                         /*num_per_page=*/1);
568 
569   // Add the first three states. Remember, the first page for each result state
570   // won't be cached (since it is returned immediately from RankAndPaginate).
571   // Each result state has a page size of 1 and a result set of 2 hits. So each
572   // result will take up one hit of our three hit budget.
573   ResultStateManager result_state_manager(/*max_total_hits=*/3,
574                                           document_store());
575   ICING_ASSERT_OK_AND_ASSIGN(
576       PageResultState page_result_state1,
577       result_state_manager.RankAndPaginate(std::move(result_state1)));
578   ICING_ASSERT_OK_AND_ASSIGN(
579       PageResultState page_result_state2,
580       result_state_manager.RankAndPaginate(std::move(result_state2)));
581   ICING_ASSERT_OK_AND_ASSIGN(
582       PageResultState page_result_state3,
583       result_state_manager.RankAndPaginate(std::move(result_state3)));
584 
585   // GetNextPage for result state 1 should return its result and decrement the
586   // number of cached hits to 2.
587   ICING_ASSERT_OK_AND_ASSIGN(
588       page_result_state1,
589       result_state_manager.GetNextPage(page_result_state1.next_page_token));
590   EXPECT_THAT(page_result_state1.scored_document_hits,
591               ElementsAre(EqualsScoredDocumentHit(CreateScoredHit(
592                   /*document_id=*/0))));
593 
594   // If retrieving the next page for result state 1 correctly decremented the
595   // current hit count to 2, then adding state 4 should still be within our
596   // budget and no other result states should be evicted.
597   ResultState result_state4 =
598       CreateResultState({AddScoredDocument(/*document_id=*/6),
599                          AddScoredDocument(/*document_id=*/7)},
600                         /*num_per_page=*/1);
601   ICING_ASSERT_OK_AND_ASSIGN(
602       PageResultState page_result_state4,
603       result_state_manager.RankAndPaginate(std::move(result_state4)));
604 
605   EXPECT_THAT(
606       result_state_manager.GetNextPage(page_result_state1.next_page_token),
607       StatusIs(libtextclassifier3::StatusCode::NOT_FOUND));
608 
609   ICING_ASSERT_OK_AND_ASSIGN(
610       page_result_state2,
611       result_state_manager.GetNextPage(page_result_state2.next_page_token));
612   EXPECT_THAT(page_result_state2.scored_document_hits,
613               ElementsAre(EqualsScoredDocumentHit(CreateScoredHit(
614                   /*document_id=*/2))));
615 
616   ICING_ASSERT_OK_AND_ASSIGN(
617       page_result_state3,
618       result_state_manager.GetNextPage(page_result_state3.next_page_token));
619   EXPECT_THAT(page_result_state3.scored_document_hits,
620               ElementsAre(EqualsScoredDocumentHit(CreateScoredHit(
621                   /*document_id=*/4))));
622 
623   ICING_ASSERT_OK_AND_ASSIGN(
624       page_result_state4,
625       result_state_manager.GetNextPage(page_result_state4.next_page_token));
626   EXPECT_THAT(page_result_state4.scored_document_hits,
627               ElementsAre(EqualsScoredDocumentHit(CreateScoredHit(
628                   /*document_id=*/6))));
629 }
630 
TEST_F(ResultStateManagerTest,GetNextPageShouldDecreaseCurrentHitsCountByExactlyOnePage)631 TEST_F(ResultStateManagerTest,
632        GetNextPageShouldDecreaseCurrentHitsCountByExactlyOnePage) {
633   ResultState result_state1 =
634       CreateResultState({AddScoredDocument(/*document_id=*/0),
635                          AddScoredDocument(/*document_id=*/1)},
636                         /*num_per_page=*/1);
637   ResultState result_state2 =
638       CreateResultState({AddScoredDocument(/*document_id=*/2),
639                          AddScoredDocument(/*document_id=*/3)},
640                         /*num_per_page=*/1);
641   ResultState result_state3 =
642       CreateResultState({AddScoredDocument(/*document_id=*/4),
643                          AddScoredDocument(/*document_id=*/5)},
644                         /*num_per_page=*/1);
645 
646   // Add the first three states. Remember, the first page for each result state
647   // won't be cached (since it is returned immediately from RankAndPaginate).
648   // Each result state has a page size of 1 and a result set of 2 hits. So each
649   // result will take up one hit of our three hit budget.
650   ResultStateManager result_state_manager(/*max_total_hits=*/3,
651                                           document_store());
652   ICING_ASSERT_OK_AND_ASSIGN(
653       PageResultState page_result_state1,
654       result_state_manager.RankAndPaginate(std::move(result_state1)));
655   ICING_ASSERT_OK_AND_ASSIGN(
656       PageResultState page_result_state2,
657       result_state_manager.RankAndPaginate(std::move(result_state2)));
658   ICING_ASSERT_OK_AND_ASSIGN(
659       PageResultState page_result_state3,
660       result_state_manager.RankAndPaginate(std::move(result_state3)));
661 
662   // GetNextPage for result state 1 should return its result and decrement the
663   // number of cached hits to 2.
664   ICING_ASSERT_OK_AND_ASSIGN(
665       page_result_state1,
666       result_state_manager.GetNextPage(page_result_state1.next_page_token));
667   EXPECT_THAT(page_result_state1.scored_document_hits,
668               ElementsAre(EqualsScoredDocumentHit(CreateScoredHit(
669                   /*document_id=*/0))));
670 
671   // If retrieving the next page for result state 1 correctly decremented the
672   // current hit count to 2, then adding state 4 should still be within our
673   // budget and no other result states should be evicted.
674   ResultState result_state4 =
675       CreateResultState({AddScoredDocument(/*document_id=*/6),
676                          AddScoredDocument(/*document_id=*/7)},
677                         /*num_per_page=*/1);
678   ICING_ASSERT_OK_AND_ASSIGN(
679       PageResultState page_result_state4,
680       result_state_manager.RankAndPaginate(std::move(result_state4)));
681 
682   // If retrieving the next page for result state 1 correctly decremented the
683   // current hit count to 2 and adding state 4 correctly incremented it to 3,
684   // then adding this result state should trigger the eviction of state 2.
685   ResultState result_state5 =
686       CreateResultState({AddScoredDocument(/*document_id=*/8),
687                          AddScoredDocument(/*document_id=*/9)},
688                         /*num_per_page=*/1);
689   ICING_ASSERT_OK_AND_ASSIGN(
690       PageResultState page_result_state5,
691       result_state_manager.RankAndPaginate(std::move(result_state5)));
692 
693   EXPECT_THAT(
694       result_state_manager.GetNextPage(page_result_state1.next_page_token),
695       StatusIs(libtextclassifier3::StatusCode::NOT_FOUND));
696 
697   EXPECT_THAT(
698       result_state_manager.GetNextPage(page_result_state2.next_page_token),
699       StatusIs(libtextclassifier3::StatusCode::NOT_FOUND));
700 
701   ICING_ASSERT_OK_AND_ASSIGN(
702       page_result_state3,
703       result_state_manager.GetNextPage(page_result_state3.next_page_token));
704   EXPECT_THAT(page_result_state3.scored_document_hits,
705               ElementsAre(EqualsScoredDocumentHit(CreateScoredHit(
706                   /*document_id=*/4))));
707 
708   ICING_ASSERT_OK_AND_ASSIGN(
709       page_result_state4,
710       result_state_manager.GetNextPage(page_result_state4.next_page_token));
711   EXPECT_THAT(page_result_state4.scored_document_hits,
712               ElementsAre(EqualsScoredDocumentHit(CreateScoredHit(
713                   /*document_id=*/6))));
714 
715   ICING_ASSERT_OK_AND_ASSIGN(
716       page_result_state5,
717       result_state_manager.GetNextPage(page_result_state5.next_page_token));
718   EXPECT_THAT(page_result_state5.scored_document_hits,
719               ElementsAre(EqualsScoredDocumentHit(CreateScoredHit(
720                   /*document_id=*/8))));
721 }
722 
TEST_F(ResultStateManagerTest,AddingOverBudgetResultStateShouldEvictAllStates)723 TEST_F(ResultStateManagerTest,
724        AddingOverBudgetResultStateShouldEvictAllStates) {
725   ResultState result_state1 =
726       CreateResultState({AddScoredDocument(/*document_id=*/0),
727                          AddScoredDocument(/*document_id=*/1),
728                          AddScoredDocument(/*document_id=*/2)},
729                         /*num_per_page=*/1);
730   ResultState result_state2 =
731       CreateResultState({AddScoredDocument(/*document_id=*/3),
732                          AddScoredDocument(/*document_id=*/4)},
733                         /*num_per_page=*/1);
734 
735   // Add the first two states. Remember, the first page for each result state
736   // won't be cached (since it is returned immediately from RankAndPaginate).
737   // Each result state has a page size of 1. So 3 hits will remain cached.
738   ResultStateManager result_state_manager(/*max_total_hits=*/4,
739                                           document_store());
740   ICING_ASSERT_OK_AND_ASSIGN(
741       PageResultState page_result_state1,
742       result_state_manager.RankAndPaginate(std::move(result_state1)));
743   ICING_ASSERT_OK_AND_ASSIGN(
744       PageResultState page_result_state2,
745       result_state_manager.RankAndPaginate(std::move(result_state2)));
746 
747   // Add a result state that is larger than the entire budget. This should
748   // result in all previous result states being evicted, the first hit from
749   // result state 3 being returned and the next four hits being cached (the last
750   // hit should be dropped because it exceeds the max).
751   ResultState result_state3 =
752       CreateResultState({AddScoredDocument(/*document_id=*/5),
753                          AddScoredDocument(/*document_id=*/6),
754                          AddScoredDocument(/*document_id=*/7),
755                          AddScoredDocument(/*document_id=*/8),
756                          AddScoredDocument(/*document_id=*/9),
757                          AddScoredDocument(/*document_id=*/10)},
758                         /*num_per_page=*/1);
759   ICING_ASSERT_OK_AND_ASSIGN(
760       PageResultState page_result_state3,
761       result_state_manager.RankAndPaginate(std::move(result_state3)));
762 
763   // GetNextPage for result state 1 and 2 should return NOT_FOUND.
764   EXPECT_THAT(
765       result_state_manager.GetNextPage(page_result_state1.next_page_token),
766       StatusIs(libtextclassifier3::StatusCode::NOT_FOUND));
767 
768   EXPECT_THAT(
769       result_state_manager.GetNextPage(page_result_state2.next_page_token),
770       StatusIs(libtextclassifier3::StatusCode::NOT_FOUND));
771 
772   // Only the next four results in state 3 should be retrievable.
773   ICING_ASSERT_OK_AND_ASSIGN(
774       page_result_state3,
775       result_state_manager.GetNextPage(page_result_state3.next_page_token));
776   EXPECT_THAT(page_result_state3.scored_document_hits,
777               ElementsAre(EqualsScoredDocumentHit(CreateScoredHit(
778                   /*document_id=*/9))));
779 
780   ICING_ASSERT_OK_AND_ASSIGN(
781       page_result_state3,
782       result_state_manager.GetNextPage(page_result_state3.next_page_token));
783   EXPECT_THAT(page_result_state3.scored_document_hits,
784               ElementsAre(EqualsScoredDocumentHit(CreateScoredHit(
785                   /*document_id=*/8))));
786 
787   ICING_ASSERT_OK_AND_ASSIGN(
788       page_result_state3,
789       result_state_manager.GetNextPage(page_result_state3.next_page_token));
790   EXPECT_THAT(page_result_state3.scored_document_hits,
791               ElementsAre(EqualsScoredDocumentHit(CreateScoredHit(
792                   /*document_id=*/7))));
793 
794   ICING_ASSERT_OK_AND_ASSIGN(
795       page_result_state3,
796       result_state_manager.GetNextPage(page_result_state3.next_page_token));
797   EXPECT_THAT(page_result_state3.scored_document_hits,
798               ElementsAre(EqualsScoredDocumentHit(CreateScoredHit(
799                   /*document_id=*/6))));
800 
801   // The final result should have been dropped because it exceeded the budget.
802   EXPECT_THAT(
803       result_state_manager.GetNextPage(page_result_state3.next_page_token),
804       StatusIs(libtextclassifier3::StatusCode::NOT_FOUND));
805 }
806 
TEST_F(ResultStateManagerTest,AddingResultStateShouldEvictOverBudgetResultState)807 TEST_F(ResultStateManagerTest,
808        AddingResultStateShouldEvictOverBudgetResultState) {
809   ResultStateManager result_state_manager(/*max_total_hits=*/4,
810                                           document_store());
811   // Add a result state that is larger than the entire budget. The entire result
812   // state will still be cached
813   ResultState result_state1 =
814       CreateResultState({AddScoredDocument(/*document_id=*/0),
815                          AddScoredDocument(/*document_id=*/1),
816                          AddScoredDocument(/*document_id=*/2),
817                          AddScoredDocument(/*document_id=*/3),
818                          AddScoredDocument(/*document_id=*/4),
819                          AddScoredDocument(/*document_id=*/5)},
820                         /*num_per_page=*/1);
821   ICING_ASSERT_OK_AND_ASSIGN(
822       PageResultState page_result_state1,
823       result_state_manager.RankAndPaginate(std::move(result_state1)));
824 
825   // Add a result state. Because state2 + state1 is larger than the budget,
826   // state1 should be evicted.
827   ResultState result_state2 =
828       CreateResultState({AddScoredDocument(/*document_id=*/6),
829                          AddScoredDocument(/*document_id=*/7)},
830                         /*num_per_page=*/1);
831   ICING_ASSERT_OK_AND_ASSIGN(
832       PageResultState page_result_state2,
833       result_state_manager.RankAndPaginate(std::move(result_state2)));
834 
835   // state1 should have been evicted and state2 should still be retrievable.
836   EXPECT_THAT(
837       result_state_manager.GetNextPage(page_result_state1.next_page_token),
838       StatusIs(libtextclassifier3::StatusCode::NOT_FOUND));
839 
840   ICING_ASSERT_OK_AND_ASSIGN(
841       page_result_state2,
842       result_state_manager.GetNextPage(page_result_state2.next_page_token));
843   EXPECT_THAT(page_result_state2.scored_document_hits,
844               ElementsAre(EqualsScoredDocumentHit(CreateScoredHit(
845                   /*document_id=*/6))));
846 }
847 
TEST_F(ResultStateManagerTest,ShouldGetSnippetContext)848 TEST_F(ResultStateManagerTest, ShouldGetSnippetContext) {
849   ResultSpecProto result_spec = CreateResultSpec(/*num_per_page=*/1);
850   result_spec.mutable_snippet_spec()->set_num_to_snippet(5);
851   result_spec.mutable_snippet_spec()->set_num_matches_per_property(5);
852   result_spec.mutable_snippet_spec()->set_max_window_utf32_length(5);
853 
854   SearchSpecProto search_spec;
855   search_spec.set_term_match_type(TermMatchType::EXACT_ONLY);
856 
857   SectionRestrictQueryTermsMap query_terms_map;
858   query_terms_map.emplace("term1", std::unordered_set<std::string>());
859 
860   ResultState original_result_state = ResultState(
861       /*scored_document_hits=*/{AddScoredDocument(/*document_id=*/0),
862                                 AddScoredDocument(/*document_id=*/1)},
863       query_terms_map, search_spec, CreateScoringSpec(), result_spec,
864       document_store());
865 
866   ResultStateManager result_state_manager(
867       /*max_total_hits=*/std::numeric_limits<int>::max(), document_store());
868   ICING_ASSERT_OK_AND_ASSIGN(
869       PageResultState page_result_state,
870       result_state_manager.RankAndPaginate(std::move(original_result_state)));
871 
872   ASSERT_THAT(page_result_state.next_page_token, Gt(kInvalidNextPageToken));
873 
874   EXPECT_THAT(page_result_state.snippet_context.match_type,
875               Eq(TermMatchType::EXACT_ONLY));
876   EXPECT_TRUE(page_result_state.snippet_context.query_terms.find("term1") !=
877               page_result_state.snippet_context.query_terms.end());
878   EXPECT_THAT(page_result_state.snippet_context.snippet_spec,
879               EqualsProto(result_spec.snippet_spec()));
880 }
881 
TEST_F(ResultStateManagerTest,ShouldGetDefaultSnippetContext)882 TEST_F(ResultStateManagerTest, ShouldGetDefaultSnippetContext) {
883   ResultSpecProto result_spec = CreateResultSpec(/*num_per_page=*/1);
884   // 0 indicates no snippeting
885   result_spec.mutable_snippet_spec()->set_num_to_snippet(0);
886   result_spec.mutable_snippet_spec()->set_num_matches_per_property(0);
887   result_spec.mutable_snippet_spec()->set_max_window_utf32_length(0);
888 
889   SearchSpecProto search_spec;
890   search_spec.set_term_match_type(TermMatchType::EXACT_ONLY);
891 
892   SectionRestrictQueryTermsMap query_terms_map;
893   query_terms_map.emplace("term1", std::unordered_set<std::string>());
894 
895   ResultState original_result_state = ResultState(
896       /*scored_document_hits=*/{AddScoredDocument(/*document_id=*/0),
897                                 AddScoredDocument(/*document_id=*/1)},
898       query_terms_map, search_spec, CreateScoringSpec(), result_spec,
899       document_store());
900 
901   ResultStateManager result_state_manager(
902       /*max_total_hits=*/std::numeric_limits<int>::max(), document_store());
903   ICING_ASSERT_OK_AND_ASSIGN(
904       PageResultState page_result_state,
905       result_state_manager.RankAndPaginate(std::move(original_result_state)));
906 
907   ASSERT_THAT(page_result_state.next_page_token, Gt(kInvalidNextPageToken));
908 
909   EXPECT_THAT(page_result_state.snippet_context.query_terms, IsEmpty());
910   EXPECT_THAT(
911       page_result_state.snippet_context.snippet_spec,
912       EqualsProto(ResultSpecProto::SnippetSpecProto::default_instance()));
913   EXPECT_THAT(page_result_state.snippet_context.match_type,
914               Eq(TermMatchType::UNKNOWN));
915 }
916 
TEST_F(ResultStateManagerTest,ShouldGetCorrectNumPreviouslyReturned)917 TEST_F(ResultStateManagerTest, ShouldGetCorrectNumPreviouslyReturned) {
918   ResultState original_result_state =
919       CreateResultState({AddScoredDocument(/*document_id=*/0),
920                          AddScoredDocument(/*document_id=*/1),
921                          AddScoredDocument(/*document_id=*/2),
922                          AddScoredDocument(/*document_id=*/3),
923                          AddScoredDocument(/*document_id=*/4)},
924                         /*num_per_page=*/2);
925 
926   ResultStateManager result_state_manager(
927       /*max_total_hits=*/std::numeric_limits<int>::max(), document_store());
928 
929   // First page, 2 results
930   ICING_ASSERT_OK_AND_ASSIGN(
931       PageResultState page_result_state1,
932       result_state_manager.RankAndPaginate(std::move(original_result_state)));
933   ASSERT_THAT(page_result_state1.scored_document_hits.size(), Eq(2));
934 
935   // No previously returned results
936   EXPECT_THAT(page_result_state1.num_previously_returned, Eq(0));
937 
938   uint64_t next_page_token = page_result_state1.next_page_token;
939 
940   // Second page, 2 results
941   ICING_ASSERT_OK_AND_ASSIGN(PageResultState page_result_state2,
942                              result_state_manager.GetNextPage(next_page_token));
943   ASSERT_THAT(page_result_state2.scored_document_hits.size(), Eq(2));
944 
945   // num_previously_returned = size of first page
946   EXPECT_THAT(page_result_state2.num_previously_returned, Eq(2));
947 
948   // Third page, 1 result
949   ICING_ASSERT_OK_AND_ASSIGN(PageResultState page_result_state3,
950                              result_state_manager.GetNextPage(next_page_token));
951   ASSERT_THAT(page_result_state3.scored_document_hits.size(), Eq(1));
952 
953   // num_previously_returned = size of first and second pages
954   EXPECT_THAT(page_result_state3.num_previously_returned, Eq(4));
955 
956   // No more results
957   EXPECT_THAT(result_state_manager.GetNextPage(next_page_token),
958               StatusIs(libtextclassifier3::StatusCode::NOT_FOUND));
959 }
960 
TEST_F(ResultStateManagerTest,ShouldStoreAllHits)961 TEST_F(ResultStateManagerTest, ShouldStoreAllHits) {
962   ScoredDocumentHit scored_hit_1 = AddScoredDocument(/*document_id=*/0);
963   ScoredDocumentHit scored_hit_2 = AddScoredDocument(/*document_id=*/1);
964   ScoredDocumentHit scored_hit_3 = AddScoredDocument(/*document_id=*/2);
965   ScoredDocumentHit scored_hit_4 = AddScoredDocument(/*document_id=*/3);
966   ScoredDocumentHit scored_hit_5 = AddScoredDocument(/*document_id=*/4);
967 
968   ResultState original_result_state = CreateResultState(
969       {scored_hit_1, scored_hit_2, scored_hit_3, scored_hit_4, scored_hit_5},
970       /*num_per_page=*/2);
971 
972   ResultStateManager result_state_manager(/*max_total_hits=*/4,
973                                           document_store());
974 
975   // The 5 input scored document hits will not be truncated. The first page of
976   // two hits will be returned immediately and the other three hits will fit
977   // within our caching budget.
978 
979   // First page, 2 results
980   ICING_ASSERT_OK_AND_ASSIGN(
981       PageResultState page_result_state1,
982       result_state_manager.RankAndPaginate(std::move(original_result_state)));
983   EXPECT_THAT(page_result_state1.scored_document_hits,
984               ElementsAre(EqualsScoredDocumentHit(scored_hit_5),
985                           EqualsScoredDocumentHit(scored_hit_4)));
986 
987   uint64_t next_page_token = page_result_state1.next_page_token;
988 
989   // Second page, 2 results.
990   ICING_ASSERT_OK_AND_ASSIGN(PageResultState page_result_state2,
991                              result_state_manager.GetNextPage(next_page_token));
992   EXPECT_THAT(page_result_state2.scored_document_hits,
993               ElementsAre(EqualsScoredDocumentHit(scored_hit_3),
994                           EqualsScoredDocumentHit(scored_hit_2)));
995 
996   // Third page, 1 result.
997   ICING_ASSERT_OK_AND_ASSIGN(PageResultState page_result_state3,
998                              result_state_manager.GetNextPage(next_page_token));
999   EXPECT_THAT(page_result_state3.scored_document_hits,
1000               ElementsAre(EqualsScoredDocumentHit(scored_hit_1)));
1001 
1002   // Fourth page, 0 results.
1003   EXPECT_THAT(result_state_manager.GetNextPage(next_page_token),
1004               StatusIs(libtextclassifier3::StatusCode::NOT_FOUND));
1005 }
1006 
1007 }  // namespace
1008 }  // namespace lib
1009 }  // namespace icing
1010