enhance(test-benchmark): tweaks

Author: spencer-tb

packages/testing/src/execution_testing/cli/pytest_commands/plugins/shared/benchmarking.py

@@ -1,7 +1,7 @@
 """The module contains the pytest hooks for the gas benchmark values."""
 
-import importlib.util
 import sys
+from importlib.util import module_from_spec, spec_from_file_location
 from pathlib import Path
 
 import pytest
@@ -50,29 +50,40 @@ def pytest_configure(config: pytest.Config) -> None:
         config.op_mode = OpMode.BENCHMARKING  # type: ignore[attr-defined]
 
 
-def _load_opcode_counts_config(config: pytest.Config) -> dict[str, list[int]] | None:
+def load_opcode_counts_config(
+    config: pytest.Config,
+) -> dict[str, list[int]] | None:
     """
     Load the opcode counts configuration from `tests/benchmark/configs/fixed_opcode_counts.py`.
 
     Returns dictionary mapping scenario IDs to opcode counts, or None if config not found.
     """
-    config_path = Path(config.rootpath) / "tests" / "benchmark" / "configs" / "fixed_opcode_counts.py"
+    config_path = (
+        Path(config.rootpath)
+        / "tests"
+        / "benchmark"
+        / "configs"
+        / "fixed_opcode_counts.py"
+    )
 
     if not config_path.exists():
         return None
 
     # Load the module dynamically
-    spec = importlib.util.spec_from_file_location("fixed_opcode_counts", config_path)
+    spec = spec_from_file_location("fixed_opcode_counts", config_path)
     if spec is None or spec.loader is None:
         return None
 
-    module = importlib.util.module_from_spec(spec)
+    module = module_from_spec(spec)
     sys.modules["fixed_opcode_counts"] = module
     spec.loader.exec_module(module)
 
     return {
-        "scenario_counts": getattr(module, "SCENARIO_OPCODE_COUNTS", {}),
-        "default_counts": getattr(module, "DEFAULT_OPCODE_COUNTS", [100, 500, 1000]),
+        "scenario_configs": getattr(module, "SCENARIO_CONFIGS", {}),
+        "default_counts": getattr(module, "DEFAULT_OPCODE_COUNTS", [100]),
+        "get_opcode_counts": getattr(
+            module, "get_opcode_counts_for_scenario", None
+        ),
     }
 
 
@@ -88,7 +99,7 @@ def pytest_collection_modifyitems(
 
     # Load config data if flag provided without value (empty string)
     if fixed_opcode_count == "":
-        config_data = _load_opcode_counts_config(config)
+        config_data = load_opcode_counts_config(config)
         if config_data:
             config._opcode_counts_config = config_data  # type: ignore[attr-defined]
 
@@ -161,14 +172,34 @@ def pytest_generate_tests(metafunc: pytest.Metafunc) -> None:
                 ]
             elif fixed_opcode_counts_cli == "":
                 # Flag provided without value - load from config file
+                # Check if config data was already loaded in pytest_collection_modifyitems
                 config_data = getattr(
                     metafunc.config, "_opcode_counts_config", None
                 )
+
+                # If not loaded yet (pytest_generate_tests runs first), load it now
+                if config_data is None:
+                    config_data = load_opcode_counts_config(metafunc.config)
+                    if config_data:
+                        metafunc.config._opcode_counts_config = config_data  # type: ignore[attr-defined]
+
                 if config_data:
-                    default_counts = config_data.get(
-                        "default_counts", [100]  # Default to 100k for all.
-                    )
-                    opcode_counts_to_use = default_counts
+                    # Use the config's lookup function to get opcode counts
+                    test_name = metafunc.function.__name__
+                    get_counts_fn = config_data.get("get_opcode_counts")
+
+                    if get_counts_fn:
+                        # At this stage, we only have the test name, not full parametrized IDs
+                        # The regex patterns in the config match against test_name
+                        # e.g., "test_selfbalance.*" will match "test_selfbalance"
+                        opcode_counts_to_use = get_counts_fn(
+                            test_name, test_name
+                        )
+                    else:
+                        # Fallback to default if function not available
+                        opcode_counts_to_use = config_data.get(
+                            "default_counts", [100]
+                        )
 
             # Parametrize if we have counts to use
             if opcode_counts_to_use: