Benchmark Fixture fixes

Update plugin.py

Author: KRRT7

codeflash/benchmarking/plugin/plugin.py

@@ -16,7 +16,12 @@
 if TYPE_CHECKING:
     from codeflash.models.models import BenchmarkKey
 
-IS_PYTEST_BENCHMARK_INSTALLED = importlib.util.find_spec("pytest_benchmark") is not None
+try:
+    import pytest_benchmark
+
+    PYTEST_BENCHMARK_INSTALLED = True
+except ImportError:
+    PYTEST_BENCHMARK_INSTALLED = False
 
 
 class CodeFlashBenchmarkPlugin:
@@ -288,7 +293,7 @@ def pytest_configure(config: pytest.Config) -> None:
     """Register the benchmark marker and disable conflicting plugins."""
     config.addinivalue_line("markers", "benchmark: mark test as a benchmark that should be run with codeflash tracing")
 
-    if config.getoption("--codeflash-trace") and IS_PYTEST_BENCHMARK_INSTALLED:
+    if config.getoption("--codeflash-trace") and PYTEST_BENCHMARK_INSTALLED:
         config.option.benchmark_disable = True
         config.pluginmanager.set_blocked("pytest_benchmark")
         config.pluginmanager.set_blocked("pytest-benchmark")
@@ -302,6 +307,35 @@ def pytest_addoption(parser: pytest.Parser) -> None:
 
 @pytest.fixture
 def benchmark(request: pytest.FixtureRequest) -> object:
-    if not request.config.getoption("--codeflash-trace"):
+    """Benchmark fixture that works with or without pytest-benchmark installed."""
+    config = request.config
+
+    # If --codeflash-trace is enabled, use our implementation
+    if config.getoption("--codeflash-trace"):
+        return codeflash_benchmark_plugin.Benchmark(request)
+
+    # If pytest-benchmark is installed and --codeflash-trace is not enabled,
+    # return the normal pytest-benchmark fixture
+    if PYTEST_BENCHMARK_INSTALLED:
+        from pytest_benchmark.fixture import BenchmarkFixture as BSF  # noqa: N814
+
+        bs = getattr(config, "_benchmarksession", None)
+        if bs and bs.skip:
+            pytest.skip("Benchmarks are skipped (--benchmark-skip was used).")
+
+        node = request.node
+        marker = node.get_closest_marker("benchmark")
+        options = dict(marker.kwargs) if marker else {}
+
+        if bs:
+            return BSF(
+                node,
+                add_stats=bs.benchmarks.append,
+                logger=bs.logger,
+                warner=request.node.warn,
+                disabled=bs.disabled,
+                **dict(bs.options, **options),
+            )
         return lambda func, *args, **kwargs: func(*args, **kwargs)
-    return codeflash_benchmark_plugin.Benchmark(request)
+    
+    return lambda func, *args, **kwargs: func(*args, **kwargs)