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.
In your test, you can patch fetch_price and make the mock respond differently for different product IDs.
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.
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.
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_effectwhen 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.

