| # Copyright 2015 Google LLC |
| # |
| # Licensed under the Apache License, Version 2.0 (the "License"); |
| # you may not use this file except in compliance with the License. |
| # You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, software |
| # distributed under the License is distributed on an "AS IS" BASIS, |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| # See the License for the specific language governing permissions and |
| # limitations under the License. |
| |
| import types |
| |
| import mock |
| import pytest |
| import six |
| |
| from google.api_core import page_iterator |
| |
| |
| def test__do_nothing_page_start(): |
| assert page_iterator._do_nothing_page_start(None, None, None) is None |
| |
| |
| class TestPage(object): |
| def test_constructor(self): |
| parent = mock.sentinel.parent |
| item_to_value = mock.sentinel.item_to_value |
| |
| page = page_iterator.Page(parent, (1, 2, 3), item_to_value) |
| |
| assert page.num_items == 3 |
| assert page.remaining == 3 |
| assert page._parent is parent |
| assert page._item_to_value is item_to_value |
| assert page.raw_page is None |
| |
| def test___iter__(self): |
| page = page_iterator.Page(None, (), None, None) |
| assert iter(page) is page |
| |
| def test_iterator_calls_parent_item_to_value(self): |
| parent = mock.sentinel.parent |
| |
| item_to_value = mock.Mock( |
| side_effect=lambda iterator, value: value, spec=["__call__"] |
| ) |
| |
| page = page_iterator.Page(parent, (10, 11, 12), item_to_value) |
| page._remaining = 100 |
| |
| assert item_to_value.call_count == 0 |
| assert page.remaining == 100 |
| |
| assert six.next(page) == 10 |
| assert item_to_value.call_count == 1 |
| item_to_value.assert_called_with(parent, 10) |
| assert page.remaining == 99 |
| |
| assert six.next(page) == 11 |
| assert item_to_value.call_count == 2 |
| item_to_value.assert_called_with(parent, 11) |
| assert page.remaining == 98 |
| |
| assert six.next(page) == 12 |
| assert item_to_value.call_count == 3 |
| item_to_value.assert_called_with(parent, 12) |
| assert page.remaining == 97 |
| |
| def test_raw_page(self): |
| parent = mock.sentinel.parent |
| item_to_value = mock.sentinel.item_to_value |
| |
| raw_page = mock.sentinel.raw_page |
| |
| page = page_iterator.Page(parent, (1, 2, 3), item_to_value, raw_page=raw_page) |
| assert page.raw_page is raw_page |
| |
| with pytest.raises(AttributeError): |
| page.raw_page = None |
| |
| |
| class PageIteratorImpl(page_iterator.Iterator): |
| def _next_page(self): |
| return mock.create_autospec(page_iterator.Page, instance=True) |
| |
| |
| class TestIterator(object): |
| def test_constructor(self): |
| client = mock.sentinel.client |
| item_to_value = mock.sentinel.item_to_value |
| token = "ab13nceor03" |
| max_results = 1337 |
| |
| iterator = PageIteratorImpl( |
| client, item_to_value, page_token=token, max_results=max_results |
| ) |
| |
| assert not iterator._started |
| assert iterator.client is client |
| assert iterator.item_to_value == item_to_value |
| assert iterator.max_results == max_results |
| # Changing attributes. |
| assert iterator.page_number == 0 |
| assert iterator.next_page_token == token |
| assert iterator.num_results == 0 |
| |
| def test_pages_property_starts(self): |
| iterator = PageIteratorImpl(None, None) |
| |
| assert not iterator._started |
| |
| assert isinstance(iterator.pages, types.GeneratorType) |
| |
| assert iterator._started |
| |
| def test_pages_property_restart(self): |
| iterator = PageIteratorImpl(None, None) |
| |
| assert iterator.pages |
| |
| # Make sure we cannot restart. |
| with pytest.raises(ValueError): |
| assert iterator.pages |
| |
| def test__page_iter_increment(self): |
| iterator = PageIteratorImpl(None, None) |
| page = page_iterator.Page( |
| iterator, ("item",), page_iterator._item_to_value_identity) |
| iterator._next_page = mock.Mock(side_effect=[page, None]) |
| |
| assert iterator.num_results == 0 |
| |
| page_iter = iterator._page_iter(increment=True) |
| next(page_iter) |
| |
| assert iterator.num_results == 1 |
| |
| def test__page_iter_no_increment(self): |
| iterator = PageIteratorImpl(None, None) |
| |
| assert iterator.num_results == 0 |
| |
| page_iter = iterator._page_iter(increment=False) |
| next(page_iter) |
| |
| # results should still be 0 after fetching a page. |
| assert iterator.num_results == 0 |
| |
| def test__items_iter(self): |
| # Items to be returned. |
| item1 = 17 |
| item2 = 100 |
| item3 = 211 |
| |
| # Make pages from mock responses |
| parent = mock.sentinel.parent |
| page1 = page_iterator.Page( |
| parent, (item1, item2), page_iterator._item_to_value_identity) |
| page2 = page_iterator.Page( |
| parent, (item3,), page_iterator._item_to_value_identity) |
| |
| iterator = PageIteratorImpl(None, None) |
| iterator._next_page = mock.Mock(side_effect=[page1, page2, None]) |
| |
| items_iter = iterator._items_iter() |
| |
| assert isinstance(items_iter, types.GeneratorType) |
| |
| # Consume items and check the state of the iterator. |
| assert iterator.num_results == 0 |
| |
| assert six.next(items_iter) == item1 |
| assert iterator.num_results == 1 |
| |
| assert six.next(items_iter) == item2 |
| assert iterator.num_results == 2 |
| |
| assert six.next(items_iter) == item3 |
| assert iterator.num_results == 3 |
| |
| with pytest.raises(StopIteration): |
| six.next(items_iter) |
| |
| def test___iter__(self): |
| iterator = PageIteratorImpl(None, None) |
| iterator._next_page = mock.Mock(side_effect=[(1, 2), (3,), None]) |
| |
| assert not iterator._started |
| |
| result = list(iterator) |
| |
| assert result == [1, 2, 3] |
| assert iterator._started |
| |
| def test___iter__restart(self): |
| iterator = PageIteratorImpl(None, None) |
| |
| iter(iterator) |
| |
| # Make sure we cannot restart. |
| with pytest.raises(ValueError): |
| iter(iterator) |
| |
| def test___iter___restart_after_page(self): |
| iterator = PageIteratorImpl(None, None) |
| |
| assert iterator.pages |
| |
| # Make sure we cannot restart after starting the page iterator |
| with pytest.raises(ValueError): |
| iter(iterator) |
| |
| |
| class TestHTTPIterator(object): |
| def test_constructor(self): |
| client = mock.sentinel.client |
| path = "/foo" |
| iterator = page_iterator.HTTPIterator( |
| client, mock.sentinel.api_request, path, mock.sentinel.item_to_value |
| ) |
| |
| assert not iterator._started |
| assert iterator.client is client |
| assert iterator.path == path |
| assert iterator.item_to_value is mock.sentinel.item_to_value |
| assert iterator._items_key == "items" |
| assert iterator.max_results is None |
| assert iterator.extra_params == {} |
| assert iterator._page_start == page_iterator._do_nothing_page_start |
| # Changing attributes. |
| assert iterator.page_number == 0 |
| assert iterator.next_page_token is None |
| assert iterator.num_results == 0 |
| |
| def test_constructor_w_extra_param_collision(self): |
| extra_params = {"pageToken": "val"} |
| |
| with pytest.raises(ValueError): |
| page_iterator.HTTPIterator( |
| mock.sentinel.client, |
| mock.sentinel.api_request, |
| mock.sentinel.path, |
| mock.sentinel.item_to_value, |
| extra_params=extra_params, |
| ) |
| |
| def test_iterate(self): |
| path = "/foo" |
| item1 = {"name": "1"} |
| item2 = {"name": "2"} |
| api_request = mock.Mock(return_value={"items": [item1, item2]}) |
| iterator = page_iterator.HTTPIterator( |
| mock.sentinel.client, |
| api_request, |
| path=path, |
| item_to_value=page_iterator._item_to_value_identity, |
| ) |
| |
| assert iterator.num_results == 0 |
| |
| items_iter = iter(iterator) |
| |
| val1 = six.next(items_iter) |
| assert val1 == item1 |
| assert iterator.num_results == 1 |
| |
| val2 = six.next(items_iter) |
| assert val2 == item2 |
| assert iterator.num_results == 2 |
| |
| with pytest.raises(StopIteration): |
| six.next(items_iter) |
| |
| api_request.assert_called_once_with(method="GET", path=path, query_params={}) |
| |
| def test__has_next_page_new(self): |
| iterator = page_iterator.HTTPIterator( |
| mock.sentinel.client, |
| mock.sentinel.api_request, |
| mock.sentinel.path, |
| mock.sentinel.item_to_value, |
| ) |
| |
| # The iterator should *always* indicate that it has a next page |
| # when created so that it can fetch the initial page. |
| assert iterator._has_next_page() |
| |
| def test__has_next_page_without_token(self): |
| iterator = page_iterator.HTTPIterator( |
| mock.sentinel.client, |
| mock.sentinel.api_request, |
| mock.sentinel.path, |
| mock.sentinel.item_to_value, |
| ) |
| |
| iterator.page_number = 1 |
| |
| # The iterator should not indicate that it has a new page if the |
| # initial page has been requested and there's no page token. |
| assert not iterator._has_next_page() |
| |
| def test__has_next_page_w_number_w_token(self): |
| iterator = page_iterator.HTTPIterator( |
| mock.sentinel.client, |
| mock.sentinel.api_request, |
| mock.sentinel.path, |
| mock.sentinel.item_to_value, |
| ) |
| |
| iterator.page_number = 1 |
| iterator.next_page_token = mock.sentinel.token |
| |
| # The iterator should indicate that it has a new page if the |
| # initial page has been requested and there's is a page token. |
| assert iterator._has_next_page() |
| |
| def test__has_next_page_w_max_results_not_done(self): |
| iterator = page_iterator.HTTPIterator( |
| mock.sentinel.client, |
| mock.sentinel.api_request, |
| mock.sentinel.path, |
| mock.sentinel.item_to_value, |
| max_results=3, |
| page_token=mock.sentinel.token, |
| ) |
| |
| iterator.page_number = 1 |
| |
| # The iterator should indicate that it has a new page if there |
| # is a page token and it has not consumed more than max_results. |
| assert iterator.num_results < iterator.max_results |
| assert iterator._has_next_page() |
| |
| def test__has_next_page_w_max_results_done(self): |
| |
| iterator = page_iterator.HTTPIterator( |
| mock.sentinel.client, |
| mock.sentinel.api_request, |
| mock.sentinel.path, |
| mock.sentinel.item_to_value, |
| max_results=3, |
| page_token=mock.sentinel.token, |
| ) |
| |
| iterator.page_number = 1 |
| iterator.num_results = 3 |
| |
| # The iterator should not indicate that it has a new page if there |
| # if it has consumed more than max_results. |
| assert iterator.num_results == iterator.max_results |
| assert not iterator._has_next_page() |
| |
| def test__get_query_params_no_token(self): |
| iterator = page_iterator.HTTPIterator( |
| mock.sentinel.client, |
| mock.sentinel.api_request, |
| mock.sentinel.path, |
| mock.sentinel.item_to_value, |
| ) |
| |
| assert iterator._get_query_params() == {} |
| |
| def test__get_query_params_w_token(self): |
| iterator = page_iterator.HTTPIterator( |
| mock.sentinel.client, |
| mock.sentinel.api_request, |
| mock.sentinel.path, |
| mock.sentinel.item_to_value, |
| ) |
| iterator.next_page_token = "token" |
| |
| assert iterator._get_query_params() == {"pageToken": iterator.next_page_token} |
| |
| def test__get_query_params_w_max_results(self): |
| max_results = 3 |
| iterator = page_iterator.HTTPIterator( |
| mock.sentinel.client, |
| mock.sentinel.api_request, |
| mock.sentinel.path, |
| mock.sentinel.item_to_value, |
| max_results=max_results, |
| ) |
| |
| iterator.num_results = 1 |
| local_max = max_results - iterator.num_results |
| |
| assert iterator._get_query_params() == {"maxResults": local_max} |
| |
| def test__get_query_params_extra_params(self): |
| extra_params = {"key": "val"} |
| iterator = page_iterator.HTTPIterator( |
| mock.sentinel.client, |
| mock.sentinel.api_request, |
| mock.sentinel.path, |
| mock.sentinel.item_to_value, |
| extra_params=extra_params, |
| ) |
| |
| assert iterator._get_query_params() == extra_params |
| |
| def test__get_next_page_response_with_post(self): |
| path = "/foo" |
| page_response = {"items": ["one", "two"]} |
| api_request = mock.Mock(return_value=page_response) |
| iterator = page_iterator.HTTPIterator( |
| mock.sentinel.client, |
| api_request, |
| path=path, |
| item_to_value=page_iterator._item_to_value_identity, |
| ) |
| iterator._HTTP_METHOD = "POST" |
| |
| response = iterator._get_next_page_response() |
| |
| assert response == page_response |
| |
| api_request.assert_called_once_with(method="POST", path=path, data={}) |
| |
| def test__get_next_page_bad_http_method(self): |
| iterator = page_iterator.HTTPIterator( |
| mock.sentinel.client, |
| mock.sentinel.api_request, |
| mock.sentinel.path, |
| mock.sentinel.item_to_value, |
| ) |
| iterator._HTTP_METHOD = "NOT-A-VERB" |
| |
| with pytest.raises(ValueError): |
| iterator._get_next_page_response() |
| |
| |
| class TestGRPCIterator(object): |
| def test_constructor(self): |
| client = mock.sentinel.client |
| items_field = "items" |
| iterator = page_iterator.GRPCIterator( |
| client, mock.sentinel.method, mock.sentinel.request, items_field |
| ) |
| |
| assert not iterator._started |
| assert iterator.client is client |
| assert iterator.max_results is None |
| assert iterator.item_to_value is page_iterator._item_to_value_identity |
| assert iterator._method == mock.sentinel.method |
| assert iterator._request == mock.sentinel.request |
| assert iterator._items_field == items_field |
| assert ( |
| iterator._request_token_field |
| == page_iterator.GRPCIterator._DEFAULT_REQUEST_TOKEN_FIELD |
| ) |
| assert ( |
| iterator._response_token_field |
| == page_iterator.GRPCIterator._DEFAULT_RESPONSE_TOKEN_FIELD |
| ) |
| # Changing attributes. |
| assert iterator.page_number == 0 |
| assert iterator.next_page_token is None |
| assert iterator.num_results == 0 |
| |
| def test_constructor_options(self): |
| client = mock.sentinel.client |
| items_field = "items" |
| request_field = "request" |
| response_field = "response" |
| iterator = page_iterator.GRPCIterator( |
| client, |
| mock.sentinel.method, |
| mock.sentinel.request, |
| items_field, |
| item_to_value=mock.sentinel.item_to_value, |
| request_token_field=request_field, |
| response_token_field=response_field, |
| max_results=42, |
| ) |
| |
| assert iterator.client is client |
| assert iterator.max_results == 42 |
| assert iterator.item_to_value is mock.sentinel.item_to_value |
| assert iterator._method == mock.sentinel.method |
| assert iterator._request == mock.sentinel.request |
| assert iterator._items_field == items_field |
| assert iterator._request_token_field == request_field |
| assert iterator._response_token_field == response_field |
| |
| def test_iterate(self): |
| request = mock.Mock(spec=["page_token"], page_token=None) |
| response1 = mock.Mock(items=["a", "b"], next_page_token="1") |
| response2 = mock.Mock(items=["c"], next_page_token="2") |
| response3 = mock.Mock(items=["d"], next_page_token="") |
| method = mock.Mock(side_effect=[response1, response2, response3]) |
| iterator = page_iterator.GRPCIterator( |
| mock.sentinel.client, method, request, "items" |
| ) |
| |
| assert iterator.num_results == 0 |
| |
| items = list(iterator) |
| assert items == ["a", "b", "c", "d"] |
| |
| method.assert_called_with(request) |
| assert method.call_count == 3 |
| assert request.page_token == "2" |
| |
| def test_iterate_with_max_results(self): |
| request = mock.Mock(spec=["page_token"], page_token=None) |
| response1 = mock.Mock(items=["a", "b"], next_page_token="1") |
| response2 = mock.Mock(items=["c"], next_page_token="2") |
| response3 = mock.Mock(items=["d"], next_page_token="") |
| method = mock.Mock(side_effect=[response1, response2, response3]) |
| iterator = page_iterator.GRPCIterator( |
| mock.sentinel.client, method, request, "items", max_results=3 |
| ) |
| |
| assert iterator.num_results == 0 |
| |
| items = list(iterator) |
| |
| assert items == ["a", "b", "c"] |
| assert iterator.num_results == 3 |
| |
| method.assert_called_with(request) |
| assert method.call_count == 2 |
| assert request.page_token == "1" |
| |
| |
| class GAXPageIterator(object): |
| """Fake object that matches gax.PageIterator""" |
| |
| def __init__(self, pages, page_token=None): |
| self._pages = iter(pages) |
| self.page_token = page_token |
| |
| def next(self): |
| return six.next(self._pages) |
| |
| __next__ = next |
| |
| |
| class TestGAXIterator(object): |
| def test_constructor(self): |
| client = mock.sentinel.client |
| token = "zzzyy78kl" |
| page_iter = GAXPageIterator((), page_token=token) |
| item_to_value = page_iterator._item_to_value_identity |
| max_results = 1337 |
| iterator = page_iterator._GAXIterator( |
| client, page_iter, item_to_value, max_results=max_results |
| ) |
| |
| assert not iterator._started |
| assert iterator.client is client |
| assert iterator.item_to_value is item_to_value |
| assert iterator.max_results == max_results |
| assert iterator._gax_page_iter is page_iter |
| # Changing attributes. |
| assert iterator.page_number == 0 |
| assert iterator.next_page_token == token |
| assert iterator.num_results == 0 |
| |
| def test__next_page(self): |
| page_items = (29, 31) |
| page_token = "2sde98ds2s0hh" |
| page_iter = GAXPageIterator([page_items], page_token=page_token) |
| iterator = page_iterator._GAXIterator( |
| mock.sentinel.client, page_iter, page_iterator._item_to_value_identity |
| ) |
| |
| page = iterator._next_page() |
| |
| assert iterator.next_page_token == page_token |
| assert isinstance(page, page_iterator.Page) |
| assert list(page) == list(page_items) |
| |
| next_page = iterator._next_page() |
| |
| assert next_page is None |