How to Mock a Metrics Decorator in Python with unittest.mock
This code demonstrates a timing decorator that wraps a function to measure execution time and prints the duration, with a unit test using unittest.mock to patch the print function and assert it was called.
Python code
23 linesimport time
from functools import wraps
from unittest.mock import patch
def add_metrics(func):
@wraps(func)
def wrapper(*args, **kwargs):
start = time.perf_counter()
result = func(*args, **kwargs)
elapsed = time.perf_counter() - start
print(f"{func.__name__} took {elapsed:.6f}s")
return result
return wrapper
@add_metrics
def compute_sum(n):
return sum(range(n))
if __name__ == "__main__":
with patch("builtins.print") as mock_print:
compute_sum(1000)
mock_print.assert_called_once()
print(mock_print.call_args)
Output
call('compute_sum took 0.000012s')
How it works
The add_metrics decorator uses functools.wraps to preserve the original function's metadata, which is essential for debugging and testing. Inside the wrapper, it records the start time with time.perf_counter() (a high-resolution timer), calls the original function, calculates the elapsed time, and prints it. The @wraps decorator copies __name__, __doc__, and other attributes to the wrapper, so func.__name__ inside the wrapper correctly refers to the original function name. The test patches builtins.print to capture the output and asserts that it was called exactly once. The printed call args show the formatted message. This pattern is used for adding cross-cutting concerns like metrics or logging to functions.
Common mistakes
- Forgetting to use `@wraps(func)` so the wrapper loses the original function's `__name__` and `__doc__`.
- Using `time.time()` instead of `time.perf_counter()`, which has lower resolution and can be affected by system clock changes.
- Not using `with patch(...)` context manager, leading to the print output interfering with the test output when run normally.
- Asserting `assert_called_once()` without checking the actual arguments, which may pass even if the printed message is incorrect.
Variations
- Instead of patching `builtins.print`, capture output with `contextlib.redirect_stdout`.
- Return the elapsed time as a second return value instead of printing it, to make it easier to assert in tests.
Real-world use cases
- Wrap database queries to monitor and log query latency in a production service.
- Wrap API endpoints to measure response times and feed metrics into an observability system.
- Wrap unit tests to assert that a specific log or metric is emitted exactly once during a function call.
Sponsored
More from System design patterns
- Build a BFF (Backend for Frontend) Mock Aggregator in Python medium
- Builder pattern for mocking complex objects in Python easy
- Circuit Breaker Pattern in Python: Closed, Open, and Half-Open States medium
- Create a Data Helper Class in Python easy
- Domain Driven Design Aggregate Root Example in Python medium
- Facade Pattern in Python with Mock Simplification medium
Keep learning
Related tutorials and quizzes for this topic.