cleanup plugin

Author: KRRT7

codeflash/benchmarking/plugin/plugin.py

@@ -211,34 +211,27 @@ def pytest_sessionfinish(self, session, exitstatus) -> None:  # noqa: ANN001, AR
         # Close the database connection
         self.close()
 
-    @staticmethod
-    def pytest_addoption(parser: pytest.Parser) -> None:
-        parser.addoption("--codeflash-trace", action="store_true", default=False, help="Enable CodeFlash tracing")
-
-    @staticmethod
-    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:
-            object.__setattr__(config.option, "benchmark_disable", True)
-            config.pluginmanager.set_blocked("pytest_benchmark")
-            config.pluginmanager.set_blocked("pytest-benchmark")
-
     @staticmethod
     def pytest_collection_modifyitems(config: pytest.Config, items: list[pytest.Item]) -> None:
-        # Skip tests that don't have the benchmark marker
+        # Skip tests that don't have the benchmark fixture
         if not config.getoption("--codeflash-trace"):
             return
 
+        skip_no_benchmark = pytest.mark.skip(reason="Test requires benchmark fixture")
         for item in items:
+            # Check for direct benchmark fixture usage
+            has_fixture = hasattr(item, "fixturenames") and "benchmark" in item.fixturenames
+
             # Check for @pytest.mark.benchmark marker
+            has_marker = False
             if hasattr(item, "get_closest_marker"):
                 marker = item.get_closest_marker("benchmark")
-                if marker is None:
-                    item.add_marker(pytest.mark.skip(reason="Test requires benchmark marker"))
+                if marker is not None:
+                    has_marker = True
+
+            # Skip if neither fixture nor marker is present
+            if not (has_fixture or has_marker):
+                item.add_marker(skip_no_benchmark)
 
     # Benchmark fixture
     class Benchmark:  # noqa: D106
@@ -289,21 +282,19 @@ def _run_benchmark(self, func, *args, **kwargs):  # noqa: ANN001, ANN002, ANN003
 
             return result
 
-    @staticmethod
-    @pytest.fixture
-    def benchmark(request: pytest.FixtureRequest) -> object:
-        """Fixture to provide the benchmark functionality."""
-        if not request.config.getoption("--codeflash-trace"):
-            return None
-        return CodeFlashBenchmarkPlugin.Benchmark(request)
-
 
 codeflash_benchmark_plugin = CodeFlashBenchmarkPlugin()
 
 
 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:
+        config.option.benchmark_disable = True
+        config.pluginmanager.set_blocked("pytest_benchmark")
+        config.pluginmanager.set_blocked("pytest-benchmark")
+
 
 def pytest_addoption(parser: pytest.Parser) -> None:
     parser.addoption(

tests/test_trace_benchmarks.py

@@ -33,7 +33,7 @@ def test_trace_benchmarks() -> None:
         function_calls = cursor.fetchall()
 
         # Assert the length of function calls
-        assert len(function_calls) == 8, f"Expected 8 function calls, but got {len(function_calls)}"
+        assert len(function_calls) == 7, f"Expected 7 function calls, but got {len(function_calls)}"
 
         bubble_sort_path = (project_root / "bubble_sort_codeflash_trace.py").as_posix()
         process_and_bubble_sort_path = (project_root / "process_and_bubble_sort_codeflash_trace.py").as_posix()