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.

Easy Python 3.9+ Aug 9, 2026 System design patterns 15 views 0 copies

Python code

23 lines
Python 3.9+
import 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

stdout
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

  1. Instead of patching `builtins.print`, capture output with `contextlib.redirect_stdout`.
  2. 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

Run this sample

Open the browser IDE to tweak the example and see results without installing anything.

Open editor

More from System design patterns

Related tutorials and quizzes for this topic.