Return Different Values for Each Call of A Mock

When unit-testing a function that make multiple external calls/requests, we want to mock the actual call and return the mock results. For each call, we might want to return different results.

This can be achieved easily with unittest.mock.Mock.side_effect. Note the property name is side_effect (singular), not side_effects!

side_effect can be a list of value, where each value corresponds to the return value for a single call
to simulate an exception, just set one of the list element to an actual Exception

Let’s look at the following code snippet:

# data_fetch.py

import requests


def fetch_product_data(product_ids: list[str]):
    service_url = "https://my.service.url"

    product_data = []

    for _id in product_ids:
        try:
            response = requests.get(f"{service_url}/product?id={_id}", timeout=5)
        except TimeoutError as e:
            print("fetch data time out")
            product_data.append(None)
            continue

        if response.status_code != 200:
            print(f"Failed to fetch data: {response.json()}")
            product_data.append(None)
            continue

        json_response = response.json()
        product_data.append(
            {
                "product_id": _id,
                "price": json_response.get("price"),
                "stock": json_response.get("stock"),
            }
        )

    return product_data

The function fetch_product_data get a list of product ids and fetch the product info from some remote service. If everything is okay, it will get some product info, otherwise, it will just use None for that product.

To unit-test this function, side_effect makes perfect sense. Here is our sample test script:

# test_data_fetch.py
from unittest import mock

from data_fetch import fetch_product_data


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

    def json(self):
        return self.json_data


def test_fetch_product_data():

    fake_product_ids = ["123", "456", "789", "hhh"]
    mock_responses = [
        MockResponse(
            {
                "price": 12.5,
                "stock": 100,
                "title": "milk xyz",
            },
            status_code=200,
        ),
        TimeoutError("Timeout fetching data"),
        MockResponse(
            {
                "price": 24.99,
                "stock": 123,
                "title": "cola",
            },
            status_code=200,
        ),
        MockResponse(
            {
                "error": "internal error",
            },
            status_code=503,
        ),
    ]

    expected_result = [
        {
            "product_id": "123",
            "price": 12.5,
            "stock": 100,
        },
        None,
        {
            "product_id": "789",
            "price": 24.99,
            "stock": 123,
        },
        None,
    ]

    with mock.patch("data_fetch.requests.get", spec_set=True) as mock_get:
        mock_get.side_effect = mock_responses

        actual_result = fetch_product_data(
            product_ids=fake_product_ids,
        )

        assert actual_result == expected_result

The above code is pretty self-explanatory. One thing to note: it is better to use spec_set=True, so that the mock object is not too wild!

spec_set: A stricter variant of spec. If used, attempting to set or get an attribute on the mock that isn’t on the object passed as spec_set will raise an AttributeError.

Otherwise, when you access an non-existent attribute of the mock object, it will be created automatically. I was tripped by this: I miss-spelled the side_effect to side_effects and the test script failed.

Related