From b30840c2bec00db5b9ecfef4a3e9c47ad68e89ae Mon Sep 17 00:00:00 2001
From: Arthur Pastel <arthur.pastel@gmail.com>
Date: Tue, 27 Jun 2023 15:32:18 +0200
Subject: [PATCH] fix: support benchmark.extra_info parameters on the fixture

---
 src/pytest_codspeed/plugin.py | 20 +++++++++++++-------
 tests/test_pytest_plugin.py   | 15 +++++++++++++++
 2 files changed, 28 insertions(+), 7 deletions(-)

diff --git a/src/pytest_codspeed/plugin.py b/src/pytest_codspeed/plugin.py
index 825a156..e48a2da 100644
--- a/src/pytest_codspeed/plugin.py
+++ b/src/pytest_codspeed/plugin.py
@@ -3,7 +3,7 @@
 import pkgutil
 import sys
 from dataclasses import dataclass, field
-from typing import TYPE_CHECKING, Any, Callable, List, Optional, Tuple, Union
+from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple, Union
 
 import pytest
 from _pytest.fixtures import FixtureManager
@@ -226,21 +226,27 @@ def pytest_runtest_protocol(item: "pytest.Item", nextitem: Union["pytest.Item",
     return reports  # Deny further protocol hooks execution
 
 
-@pytest.fixture(scope="function")
-def codspeed_benchmark(request: "pytest.FixtureRequest") -> Callable:
-    plugin = get_plugin(request.config)
+class BenchmarkFixture:
+    def __init__(self, request: "pytest.FixtureRequest"):
+        self.extra_info: Dict = {}
+
+        self._request = request
 
-    def run(func: Callable[..., Any], *args: Any, **kwargs: Any) -> Any:
+    def __call__(self, func: Callable[..., Any], *args: Any, **kwargs: Any) -> Any:
+        plugin = get_plugin(self._request.config)
         plugin.benchmark_count += 1
         if plugin.is_codspeed_enabled and plugin.should_measure:
             assert plugin.lib is not None
             _run_with_instrumentation(
-                plugin.lib, request.node.nodeid, func, *args, **kwargs
+                plugin.lib, self._request.node.nodeid, func, *args, **kwargs
             )
         else:
             func(*args, **kwargs)
 
-    return run
+
+@pytest.fixture(scope="function")
+def codspeed_benchmark(request: "pytest.FixtureRequest") -> Callable:
+    return BenchmarkFixture(request)
 
 
 if not IS_PYTEST_BENCHMARK_INSTALLED:
diff --git a/tests/test_pytest_plugin.py b/tests/test_pytest_plugin.py
index d82cebd..1503f2b 100644
--- a/tests/test_pytest_plugin.py
+++ b/tests/test_pytest_plugin.py
@@ -245,6 +245,21 @@ def _():
     )
 
 
+def test_pytest_benchmark_extra_info(pytester: pytest.Pytester) -> None:
+    """https://pytest-benchmark.readthedocs.io/en/latest/usage.html#extra-info"""
+    pytester.makepyfile(
+        """
+        import time
+
+        def test_my_stuff(benchmark):
+            benchmark.extra_info['foo'] = 'bar'
+            benchmark(time.sleep, 0.02)
+        """
+    )
+    result = pytester.runpytest("--codspeed")
+    assert result.ret == 0, "the run should have succeeded"
+
+
 @skip_without_valgrind
 @skip_without_perf_trampoline
 def test_perf_maps_generation(pytester: pytest.Pytester, codspeed_env) -> None: