Mock two separate responses to same function in same test

Question

In my previous question, I asked how to Mock a class that wraps requests.get in my class. The answer provided works well, if I am only calling requests.get once. However, it turns out my class is more complicated than I made my example.

My class calls request.get twice. Once at initialization, because it hits an API end point that returns API values I need to use in my actual request, and once when I make my .fetch call.

import requests
class ExampleAPI(object):
    def __init__(self):
        self.important_tokens = requests.get(url_to_tokens)['tokens']

    def fetch(self, url, params=None, key=None, token=None, **kwargs):
        return requests.get(url, params=self.important_tokens).json() 

Now, it turns out I need to create two mock responses. One for the initialization and one for the .fetch. Using the code from the previous answer:

@patch('mymodule.requests.get')
def test_fetch(self, fake_get):
    expected = {"result": "True"}
    fake_get.return_value.json.return_value = expected
    e = ExampleAPI()    # This needs one set of mocked responses
    self.assertEqual(e.fetch('http://my.api.url.example.com'), expected)    # This needs a second set

How can I create seperate responses for these two seperate calls to request.get?

Answer 1

You can assign an iterable to the side_effects attribute of a mock object; each time the mock is called, it returns the next item of the iterable.

fake_responses = [Mock(), Mock()]
fake_responses[0].json.return_value = ...
fake_responses[1].json.return_value = ...
fake_get.side_effects = fake_responses

Answer 2

Looks like the previous answer was using "side_effects" instead of "side_effect". This is how you can do that in Python 3:

import requests
import unittest
from unittest import mock
from unittest.mock import Mock


class Tests(unittest.TestCase):

    @mock.patch('requests.get')
    def test_post_price_band(self, fake_get):
        fake_responses = [Mock(), Mock()]
        fake_responses[0].json.return_value = {"a": 1}
        fake_responses[1].json.return_value = {"b": 2}
        fake_get.side_effect = fake_responses

        r1 = requests.get('https://www.api.com').json()
        self.assertEqual(r1, {"a": 1})

        r2 = requests.get('https://www.api.com').json()
        self.assertEqual(r2, {"b": 2})

Alternatively you could implement it like so:

class MockResponse:
    def __init__(self, json_data, status_code=requests.codes.ok):
        self.json_data = json_data
        self.status_code = status_code

    def json(self):
        return self.json_data


class Tests(unittest.TestCase):

    @mock.patch('requests.get')
    def test_post_price_band(self, fake_get):
        fake_get.side_effect = [
            MockResponse({"a": 1}),
            MockResponse({"b": 2})
        ]

        r1 = requests.get('https://www.api.com')
        self.assertEqual(r1.status_code, requests.codes.ok)
        self.assertEqual(r1.json(), {"a": 1})

        r2 = requests.get('https://www.api.com')
        self.assertEqual(r2.status_code, requests.codes.ok)
        self.assertEqual(r2.json(), {"b": 2})

Also check out this library to help you: https://github.com/getsentry/responses

Answer 3

I suppose an improvement to the given answers could be passing an object to the side_effect parameter. In case of an iterable, let's suppose we are mocking requests.get which is called n number of times, to mock every request.get call we pass an iterable of length n.

A possible improvement over an iterable could be an object which returns a mock function based on one of the parameters the mock function is supposed to use. In case of requests.get it can be the url. The following class demonstrates this use.

class UrlSideEffect:
    def __init__(self, url_fn_map: dict):
        """this class returns a function according to url being passed in the mock call"""
        self.url_fn_map = url_fn_map
    
    def __call__(self, *args, **kwargs):
        current_url = kwargs["url"]
        f = self.url_fn_map.get(current_url)
        return f(*args, **kwargs)

A Mock response class which behaves as the mocked object the mocked function will return:

from requests import Response
from requests.structures import CaseInsensitiveDict

class MockResponse(Response):
    def __init__(self, status_code: int, headers: Dict, text: str) -> None:
        super().__init__()
        self.status_code = status_code
        self.headers = headers
        self._text = text

    text = property(lambda obj: obj._text)

Putting together everything ...

mocked_response_url1 = MockResponse(200, {headers:1}, "{json_response:ok, url:url1}")
mocked_response_url2 = MockResponse(200, {headers:1}, "{json_response:ok, url:url2}")

side_effect_mocked_response = UrlSideEffect(
    url_fn_map={"https://url1/": lambda *args, **kwargs: mocked_response_url1,
                "https://url2/":lambda *args, **kwargs: mocked_response_url2
               }
          )

import unittest
class TestExampleAPI(unittest.TestCase):
    @mock.patch("requests.get", side_effect=side_effect_mocked_response)
    def test_initialization_and_fetch(self, _: Any) -> None:
        api = ExampleAPI() # requests.get will return mocked_response_url1
        resp = api.fetch(url2) #requests.get will return mocked_response_url2
        
        # here goes whatever you want to assert ... 

This approach can also be used in case of asynchronous requests. Imagine having a list of tasks being executed by asyncio.gather. Due to the order of requests not being synchronous we cannot pass an iterable to the side_effect parameter. Hence in such case the class UrlSideEffect can be useful as it is tied to url and not the order of function calls.

Please also have a look at this answer which inspired my answer.