fix test failures

Author: KRRT7

codeflash/benchmarking/plugin/plugin.py

@@ -5,24 +5,18 @@
 import sqlite3
 import sys
 import time
-from dataclasses import dataclass
 from pathlib import Path
+from typing import TYPE_CHECKING
 
 import pytest
 
 from codeflash.benchmarking.codeflash_trace import codeflash_trace
 from codeflash.code_utils.code_utils import module_name_from_file_path
 
-IS_PYTEST_BENCHMARK_INSTALLED = importlib.util.find_spec("pytest_benchmark") is not None
-
-
-@dataclass(frozen=True)
-class BenchmarkKey:
-    module_path: str
-    function_name: str
+if TYPE_CHECKING:
+    from codeflash.models.models import BenchmarkKey
 
-    def __str__(self) -> str:
-        return f"{self.module_path}::{self.function_name}"
+IS_PYTEST_BENCHMARK_INSTALLED = importlib.util.find_spec("pytest_benchmark") is not None
 
 
 class CodeFlashBenchmarkPlugin:
@@ -83,6 +77,8 @@ def close(self) -> None:
 
     @staticmethod
     def get_function_benchmark_timings(trace_path: Path) -> dict[str, dict[BenchmarkKey, int]]:
+        from codeflash.models.models import BenchmarkKey
+
         """Process the trace file and extract timing data for all functions.
 
         Args:
@@ -143,6 +139,8 @@ def get_function_benchmark_timings(trace_path: Path) -> dict[str, dict[Benchmark
 
     @staticmethod
     def get_benchmark_timings(trace_path: Path) -> dict[BenchmarkKey, int]:
+        from codeflash.models.models import BenchmarkKey
+
         """Extract total benchmark timings from trace files.
 
         Args:

codeflash/models/models.py

@@ -13,6 +13,7 @@
 import re
 import sys
 from collections.abc import Collection
+from dataclasses import dataclass as dcdataclass
 from enum import Enum, IntEnum
 from pathlib import Path
 from re import Pattern
@@ -83,7 +84,7 @@ class BestOptimization(BaseModel):
     winning_replay_benchmarking_test_results: Optional[TestResults] = None
 
 
-@dataclass(frozen=True)
+@dcdataclass(frozen=True)
 class BenchmarkKey:
     module_path: str
     function_name: str