How to Benchmark Python Code with pytest-benchmark and mocks

import time
from unittest.mock import Mock, patch

import pytest
from pytest_benchmark.fixture import BenchmarkFixture


def heavy_operation(data: list[int]) -> int:
    """Simulates a CPU-bound operation."""
    return sum(x * x for x in data)


def test_heavy_operation_benchmark(benchmark: BenchmarkFixture) -> None:
    """Benchmarks the heavy operation."""
    data = list(range(1000))
    result = benchmark(heavy_operation, data)
    assert result == sum(x * x for x in range(1000))


def test_heavy_operation_with_mock(benchmark: BenchmarkFixture) -> None:
    """Benchmark while mocking a dependency."""
    mock_logger = Mock()
    mock_logger.info.return_value = None

    def wrapped_operation(data: list[int]) -> int:
        mock_logger.info("Starting operation")
        result = heavy_operation(data)
        mock_logger.info("Finished operation")
        return result

    data = list(range(1000))
    result = benchmark(wrapped_operation, data)
    assert result == sum(x * x for x in range(1000))
    mock_logger.info.assert_called_with("Finished operation")


def test_heavy_operation_with_patch(benchmark: BenchmarkFixture) -> None:
    """Benchmark with patched time to control execution."""
    with patch("time.perf_counter", return_value=0.001):
        data = list(range(1000))
        result = benchmark(heavy_operation, data)
        assert result == sum(x * x for x in range(1000))


if __name__ == "__main__":
    # Demonstrate the functions work without pytest
    test_data = list(range(1000))
    expected = sum(x * x for x in test_data)
    assert heavy_operation(test_data) == expected
    print("All tests passed successfully")