diff --git a/google/cloud/internal/pagination_range.h b/google/cloud/internal/pagination_range.h index 6608e04ada5bb..4f65031ae8c49 100644 --- a/google/cloud/internal/pagination_range.h +++ b/google/cloud/internal/pagination_range.h @@ -93,8 +93,7 @@ class PagedStreamReader { * successful end of stream. */ typename StreamReader::result_type GetNext() { - if (current_ == page_.end()) { - if (last_page_) return Status{}; + while (current_ == page_.end() && !last_page_) { request_.set_page_token(std::move(token_)); auto response = loader_(request_); if (!response.ok()) return std::move(response).status(); @@ -102,8 +101,8 @@ class PagedStreamReader { if (token_.empty()) last_page_ = true; page_ = extractor_(*std::move(response)); current_ = page_.begin(); - if (current_ == page_.end()) return Status{}; } + if (current_ == page_.end()) return Status{}; return std::move(*current_++); } diff --git a/google/cloud/internal/pagination_range_test.cc b/google/cloud/internal/pagination_range_test.cc index 5dc1a54c7f334..7241e9134521f 100644 --- a/google/cloud/internal/pagination_range_test.cc +++ b/google/cloud/internal/pagination_range_test.cc @@ -197,6 +197,63 @@ TYPED_TEST(PaginationRangeTest, TwoPages) { EXPECT_THAT(names, ElementsAre("p1", "p2", "p3", "p4")); } +TYPED_TEST(PaginationRangeTest, EmptyAndFullPages) { + using ResponseType = TypeParam; + MockRpc mock; + EXPECT_CALL(mock, Loader) + .WillOnce([](Request const& request) { + EXPECT_EQ(CurrentOptions().get(), "EmptyAndFullPages"); + EXPECT_TRUE(request.testonly_page_token.empty()); + ResponseType response; + response.testonly_set_page_token("t1"); + return response; + }) + .WillOnce([](Request const& request) { + EXPECT_EQ(CurrentOptions().get(), "EmptyAndFullPages"); + EXPECT_EQ("t1", request.testonly_page_token); + ResponseType response; + response.testonly_set_page_token("t2"); + return response; + }) + .WillOnce([](Request const& request) { + EXPECT_EQ(CurrentOptions().get(), "EmptyAndFullPages"); + EXPECT_EQ("t2", request.testonly_page_token); + ResponseType response; + response.testonly_items.push_back(Item{"p1"}); + response.testonly_items.push_back(Item{"p2"}); + response.testonly_set_page_token("t3"); + return response; + }) + .WillOnce([](Request const& request) { + EXPECT_EQ(CurrentOptions().get(), "EmptyAndFullPages"); + EXPECT_EQ("t3", request.testonly_page_token); + ResponseType response; + response.testonly_set_page_token("t4"); + return response; + }) + .WillOnce([](Request const& request) { + EXPECT_EQ(CurrentOptions().get(), "EmptyAndFullPages"); + EXPECT_EQ("t4", request.testonly_page_token); + ResponseType response; + response.testonly_items.push_back(Item{"p3"}); + response.testonly_items.push_back(Item{"p4"}); + response.testonly_set_page_token(""); + return response; + }); + + OptionsSpan span(Options{}.set("EmptyAndFullPages")); + auto range = MakePaginationRange( + Request{}, [&mock](Request const& r) { return mock.Loader(r); }, + [](ResponseType const& r) { return r.testonly_items; }); + OptionsSpan overlay(Options{}.set("uh-oh")); + std::vector names; + for (auto& p : range) { + if (!p) break; + names.push_back(p->data); + } + EXPECT_THAT(names, ElementsAre("p1", "p2", "p3", "p4")); +} + TYPED_TEST(PaginationRangeTest, TwoPagesWithError) { using ResponseType = TypeParam; MockRpc mock;