diff --git a/appnexus/cursor.py b/appnexus/cursor.py index d3638d6..6dafd89 100644 --- a/appnexus/cursor.py +++ b/appnexus/cursor.py @@ -23,7 +23,6 @@ def __init__(self, client, service_name, representation, **specs): self.service_name = service_name self.representation = representation self.specs = specs - self.retrieved = 0 self._skip = 0 self._limit = float('inf') @@ -41,19 +40,7 @@ def __iter__(self): """Iterate over all AppNexus objects matching the specifications""" for page in self.iter_pages(): data = self.extract_data(page) - if self._skip >= len(data): - self._skip -= len(data) - continue - elif self._skip: - self._skip = 0 - data = data[self._skip:] - lasting = self._limit - self.retrieved - if not lasting: - break - elif lasting < len(data): - data = data[:lasting] for entity in data: - self.retrieved += 1 yield entity def extract_data(self, page): @@ -86,15 +73,17 @@ def get_page(self, start_element=0, num_elements=None): specs.update(start_element=start_element, num_elements=num_elements) return self.client.get(self.service_name, **specs) - def iter_pages(self, skip_elements=0): + def iter_pages(self): """Iterate as much as needed to get all available pages""" - start_element = skip_elements + start_element = self._skip + num_elements = min(self._limit, self.batch_size) count = -1 while start_element < count or count == -1: - page = self.get_page(start_element) + page = self.get_page(start_element, num_elements) yield page - start_element = page["start_element"] + page["num_elements"] - count = page["count"] + start_element = start_element + page["num_elements"] + num_elements = min(page["count"] - num_elements, self.batch_size) + count = min(page["count"], self._skip + self._limit) def count(self): """Returns the number of elements matching the specifications""" diff --git a/tests/cursor.py b/tests/cursor.py index 2ad3bda..631021a 100644 --- a/tests/cursor.py +++ b/tests/cursor.py @@ -4,7 +4,9 @@ from appnexus.client import AppNexusClient from appnexus.cursor import Cursor -from .helpers import gen_random_collection +from .helpers import gen_collection + +COLLECTION_SIZE = 324 @pytest.fixture @@ -55,7 +57,12 @@ def response_dict2(): @pytest.fixture def random_response_dict(): - return gen_random_collection(count=324) + return gen_collection(count=COLLECTION_SIZE, randomize=True) + + +@pytest.fixture +def ordered_response_dict(): + return gen_collection(count=COLLECTION_SIZE, randomize=False) @pytest.fixture @@ -71,7 +78,28 @@ def random_cursor(mocker, random_response_dict): client = AppNexusClient("test", "test") mocker.patch.object(client, "get") client.get.side_effect = random_response_dict - return Cursor(client, "campaign", representations.raw) + cursor = Cursor(client, "campaign", representations.raw) + mocker.patch.object(cursor, "get_page", wraps=cursor.get_page) + return cursor + + +@pytest.fixture +def ordered_cursor(mocker, ordered_response_dict): + client = AppNexusClient("test", "test") + mocker.patch.object(client, "get") + client.get.side_effect = ordered_response_dict + cursor = Cursor(client, "campaign", representations.raw) + mocker.patch.object(cursor, "get_page", wraps=cursor.get_page) + return cursor + + +def mock_ordered_cursor(mocker, start=0, count=COLLECTION_SIZE, factor=1): + client = AppNexusClient("test", "test") + mocker.patch.object(client, "get") + client.get.side_effect = gen_collection(start, count) * factor + cursor = Cursor(client, "campaign", representations.raw) + mocker.patch.object(cursor, "get_page", wraps=cursor.get_page) + return cursor def test_cursor_count(cursor, response_dict): @@ -164,3 +192,109 @@ def test_uncallable_representation(): def test_requests_volume_on_iteration(cursor): _ = [r for r in cursor] assert cursor.client.get.call_count == 1 + + +def test_skip_none(ordered_cursor): + results = [r for r in ordered_cursor] + assert len(results) == COLLECTION_SIZE + assert results[0]['id'] == 0 + assert results[-1]['id'] == COLLECTION_SIZE - 1 + assert ordered_cursor.get_page.call_count == 4 + + +def test_skip_ten(mocker): + skip = 10 + cursor = mock_ordered_cursor(mocker, start=skip) + cursor.skip(skip) + results = [r for r in cursor] + assert len(results) == COLLECTION_SIZE - skip + assert results[0]['id'] == skip + assert results[-1]['id'] == COLLECTION_SIZE - 1 + assert cursor.get_page.call_count == 4 + + +def test_skip_hundred_ten(mocker): + skip = 110 + cursor = mock_ordered_cursor(mocker, start=skip) + cursor.skip(skip) + results = [r for r in cursor] + assert len(results) == COLLECTION_SIZE - skip + assert results[0]['id'] == skip + assert results[-1]['id'] == COLLECTION_SIZE - 1 + assert cursor.get_page.call_count == 3 + + +def test_skip_twice(mocker): + skip = 10 + cursor = mock_ordered_cursor(mocker, start=skip, factor=2) + cursor.skip(skip) + results = [r for r in cursor] + assert len(results) == COLLECTION_SIZE - skip + assert results[0]['id'] == skip + assert cursor.get_page.call_count == 4 + results = [r for r in cursor] + assert len(results) == COLLECTION_SIZE - skip + assert results[0]['id'] == skip + assert cursor.get_page.call_count == 8 + + +def test_limit_ten(mocker): + limit = 10 + cursor = mock_ordered_cursor(mocker, count=limit) + cursor.limit(limit) + results = [r for r in cursor] + assert len(results) == limit + assert results[0]['id'] == 0 + assert results[-1]['id'] == limit - 1 + assert cursor.get_page.call_count == 1 + + +def test_limit_hundred_ten(mocker): + limit = 110 + cursor = mock_ordered_cursor(mocker, count=limit) + cursor.limit(limit) + results = [r for r in cursor] + assert len(results) == limit + assert results[0]['id'] == 0 + assert results[-1]['id'] == limit - 1 + assert cursor.get_page.call_count == 2 + + +def test_limit_thousand(mocker): + limit = 1000 + cursor = mock_ordered_cursor(mocker) + cursor.limit(limit) + results = [r for r in cursor] + assert len(results) == COLLECTION_SIZE + assert results[0]['id'] == 0 + assert results[-1]['id'] == COLLECTION_SIZE - 1 + assert cursor.get_page.call_count == 4 + + +def test_limit_twice(mocker): + limit = 50 + cursor = mock_ordered_cursor(mocker, count=limit, factor=2) + cursor.limit(limit) + results = [r for r in cursor] + assert len(results) == limit + assert results[0]['id'] == 0 + assert results[-1]['id'] == limit - 1 + assert cursor.get_page.call_count == 1 + results = [r for r in cursor] + assert len(results) == limit + assert results[0]['id'] == 0 + assert results[-1]['id'] == limit - 1 + assert cursor.get_page.call_count == 2 + + +def test_skip_and_limit(mocker): + skip = 10 + limit = 150 + cursor = mock_ordered_cursor(mocker, start=skip, count=skip + limit) + cursor.skip(skip) + cursor.limit(limit) + results = [r for r in cursor] + assert len(results) == limit + assert results[0]['id'] == skip + assert results[-1]['id'] == limit + skip - 1 + assert cursor.get_page.call_count == 2 diff --git a/tests/helpers.py b/tests/helpers.py index 7013a9a..032e7f8 100644 --- a/tests/helpers.py +++ b/tests/helpers.py @@ -1,12 +1,18 @@ import random -def gen_random_object(): - return {"id": random.randrange(1000000)} +def default_object_generator_func(index, randomize=False): + if randomize: + return {"id": random.randrange(1000000)} + return {"id": index} -def gen_random_page(num_elements=None, start_element=0, count=None, - object_type="campaigns"): +def gen_page(start_element=0, count=None, num_elements=None, + object_type="campaigns", + object_generator_func=default_object_generator_func, + randomize=False): + if not object_generator_func or not callable(object_generator_func): + raise ValueError("object_generator_func has to be set and callable") if count is None: count = random.randrange(10000) if num_elements is None: @@ -17,22 +23,40 @@ def gen_random_page(num_elements=None, start_element=0, count=None, "start_element": start_element, "num_elements": num_elements, "count": count, - object_type: [gen_random_object() for _ in range(num_elements)] + object_type: [object_generator_func(start_element + index, randomize) + for index in range(num_elements)] } -def gen_random_collection(count=None, object_type="campaigns"): +def gen_collection(start_element=0, count=None, object_type="campaigns", + object_generator_func=default_object_generator_func, + randomize=False): if count is None: - count = random.randrange(10000) + random.randrange(10000) + result = [] i = 0 - for i in range(count // 100): - random_page = gen_random_page(count=count, object_type=object_type, - start_element=i * 100) - result.append(random_page) - if count % 100 != 0: - random_page = gen_random_page(count=count, object_type=object_type, - start_element=i * 100, - num_elements=count % 100) - result.append(random_page) + volume = count - start_element + for i in range(volume // 100): + page = gen_page( + start_element=start_element + i * 100, + count=count, + object_type=object_type, + object_generator_func=object_generator_func, + randomize=randomize + ) + result.append(page) + + if volume % 100 != 0: + i = i + 1 if len(result) else 0 + page = gen_page( + start_element=start_element + i * 100, + count=count, + num_elements=volume % 100, + object_type=object_type, + object_generator_func=object_generator_func, + randomize=randomize + ) + result.append(page) + return result