'Writing a unittest for a generator method yielding successive request objects

I want to write a unittest for the __iter__ method a program of the following kind. The method returns a generator and can be used as shown in the if-block. In the example shown here, the generator yields the current page as long as the text body of the request object contains the string 'data'.

import requests


class PageGenerator:
    """Class for generating all pages for a given url."""
    def __init__(self):
        self.page = 1
        self.url = 'https://www.example.com?page={page}'
        self.res = None

    def __iter__(self):
        """Generates all available pages for the given url."""
        while self.res_has_data():
            self.url.format(page=self.page)
            self.res = requests.get(self.url)
            self.page += 1
            yield self.res

    def res_has_data(self):
        """Checks if self.res.text contains the string 'data'. If yes, proceeds to check next page, else aborts."""
        if self.res is None or 'data' in self.res.text:
            return True
        else:
            return False


if __name__ == '__main__':
    pages = PageGenerator()
    for page in pages:
        print(page.text)

Writing a unittest by using the httmock library to mock request objects seems to be straightforward, but doesn't work in the following implementation:

from unittest import TestCase
from httmock import urlmatch, HTTMock


class TestPageGenerator(TestCase):

    # Not a working test!
    def test_iter(self):
        @urlmatch(netloc=r'.*example\.com', query=r'.*page=1')
        def page_generator_mock(url, request):
            return {'content': '<html> data </html>'}

        with HTTMock(page_generator_mock):
            pages = PageGenerator()
            self.assertEqual([page.text for page in pages], ['<html> data </html>'])

The problem is that the comprehension [page.text for page in pages] returns ['<html> data </html>'] correctly in the first iteration but then proceeds to make a request to the not mocked url "https://www.example.com?page=2".

How can I write test_iter such that it tests if the __iter__ method is implemented correctly?



Sources

This article follows the attribution requirements of Stack Overflow and is licensed under CC BY-SA 3.0.

Source: Stack Overflow

Solution Source