python
mocking
unit-testing
test-driven-development
unittest-mock

Mocking python function based on input arguments

Master System Design with Codemia

Enhance your system design skills with over 120 practice problems, detailed solutions, and hands-on exercises.

Introduction

Sometimes a mock should not always return the same value. You want it to behave differently depending on the arguments it receives, just like the real function would. In Python, the standard way to do that is to use unittest.mock.patch together with a side_effect function that inspects the inputs and returns the appropriate result.

Use side_effect To Branch On Arguments

Suppose your production code calls a helper function named fetch_price.

python
1# shop.py
2
3def fetch_price(product_id):
4    raise RuntimeError("real API call not allowed in unit test")
5
6
7def total_for_items(product_ids):
8    return sum(fetch_price(product_id) for product_id in product_ids)

In your test, you can patch fetch_price and make the mock respond differently for different product IDs.

python
1import unittest
2from unittest.mock import patch
3import shop
4
5
6def fake_fetch_price(product_id):
7    if product_id == "apple":
8        return 2
9    if product_id == "banana":
10        return 3
11    raise ValueError(f"unexpected product: {product_id}")
12
13
14class ShopTests(unittest.TestCase):
15    @patch("shop.fetch_price", side_effect=fake_fetch_price)
16    def test_total_for_items(self, mock_fetch):
17        total = shop.total_for_items(["apple", "banana", "apple"])
18        self.assertEqual(total, 7)
19        self.assertEqual(mock_fetch.call_count, 3)
20
21
22if __name__ == "__main__":
23    unittest.main()

The mock now behaves like a tiny input-driven function instead of a constant stub.

A Dictionary-Based Pattern For Simple Cases

If the mapping is straightforward, a lookup table inside the side effect can make the test shorter.

python
1from unittest.mock import Mock
2
3responses = {
4    "a": 10,
5    "b": 20,
6}
7
8mock_fn = Mock(side_effect=lambda arg: responses[arg])
9
10print(mock_fn("a"))
11print(mock_fn("b"))

This works well when every known input maps directly to one output.

Returning Different Values And Raising Errors

A side effect function is not limited to return values. It can also raise exceptions for specific arguments, which is useful for testing error handling.

python
1from unittest.mock import Mock
2
3
4def fake_divide(value):
5    if value == 0:
6        raise ZeroDivisionError("division by zero")
7    return 100 / value
8
9
10mock_divide = Mock(side_effect=fake_divide)
11
12print(mock_divide(5))

That lets your test cover both successful and failing branches without needing the real dependency.

Patch The Name Where It Is Used

The biggest detail in Python mocking is patch location. Patch the symbol in the module under test, not where it was originally defined.

If shop.py imported fetch_price with from pricing import fetch_price, then the correct patch target for the test is shop.fetch_price, not pricing.fetch_price.

That rule explains many "my mock is ignored" failures. Python code uses the already-imported name in the current module namespace, so the patch must replace that name.

When A Real Function Is Better Than A Mock

If the logic is simple and deterministic, sometimes a small fake function is clearer than configuring a complex mock object with many call rules. The point of the test is readability as much as isolation.

That is why side_effect is so useful: it lets you keep mock assertions and fake behavior in one place without turning the test into a maze of conditional setup calls.

Common Pitfalls

The most common mistake is patching the wrong import path. If you patch the original definition instead of the name used by the module under test, the real function will still run.

Another issue is using return_value when the behavior depends on arguments. return_value is fixed, so input-sensitive behavior belongs in side_effect.

It is also easy to make the side effect too permissive. If unexpected inputs should fail the test, raise an exception instead of silently returning a default value.

Finally, do not over-mock pure functions that are already cheap and deterministic. Mocking is most valuable when the dependency is slow, external, or hard to control.

Summary

  • Use side_effect when a mock should behave differently for different inputs.
  • A small helper function is often the clearest way to express argument-based behavior.
  • Side effects can return values or raise exceptions depending on the input.
  • Patch the name where the code under test uses it, not necessarily where it was defined.
  • Keep mocks strict enough that unexpected arguments fail loudly.

Course illustration
Course illustration

All Rights Reserved.