perf: use set for O(1) membership check in TestFiles.add()

codeflash/models/models.py

@@ -426,9 +426,15 @@ class TestFile(BaseModel):
 
 class TestFiles(BaseModel):
     test_files: list[TestFile]
+    _seen_paths: set[Path] = PrivateAttr(default_factory=set)
+
+    def model_post_init(self, __context: Any, /) -> None:
+        self._seen_paths = {tf.instrumented_behavior_file_path for tf in self.test_files}
 
     def add(self, test_file: TestFile) -> None:
-        if test_file not in self.test_files:
+        key = test_file.instrumented_behavior_file_path
+        if key not in self._seen_paths:
+            self._seen_paths.add(key)
             self.test_files.append(test_file)
         else:
             msg = "Test file already exists in the list"

tests/test_test_files_add.py

@@ -0,0 +1,46 @@
+from pathlib import Path
+
+import pytest
+
+from codeflash.models.models import TestFile, TestFiles
+from codeflash.models.test_type import TestType
+
+
+class TestTestFilesAdd:
+    def test_add_unique_test_file(self) -> None:
+        tf = TestFiles(test_files=[])
+        test_file = TestFile(
+            instrumented_behavior_file_path=Path("/tmp/test_behavior.py"),
+            benchmarking_file_path=Path("/tmp/test_perf.py"),
+            test_type=TestType.GENERATED_REGRESSION,
+        )
+        tf.add(test_file)
+        assert len(tf.test_files) == 1
+        assert tf.test_files[0] is test_file
+
+    def test_add_duplicate_raises(self) -> None:
+        tf = TestFiles(test_files=[])
+        test_file = TestFile(
+            instrumented_behavior_file_path=Path("/tmp/test_behavior.py"),
+            benchmarking_file_path=Path("/tmp/test_perf.py"),
+            test_type=TestType.GENERATED_REGRESSION,
+        )
+        tf.add(test_file)
+        with pytest.raises(ValueError, match="Test file already exists"):
+            tf.add(test_file)
+
+    def test_add_many_files_performance(self) -> None:
+        tf = TestFiles(test_files=[])
+        for i in range(100):
+            test_file = TestFile(
+                instrumented_behavior_file_path=Path(f"/tmp/test_behavior_{i}.py"),
+                benchmarking_file_path=Path(f"/tmp/test_perf_{i}.py"),
+                test_type=TestType.GENERATED_REGRESSION,
+            )
+            tf.add(test_file)
+
+        assert len(tf.test_files) == 100
+        assert len(tf._seen_paths) == 100
+        # Verify all paths are unique in the set
+        expected_paths = {Path(f"/tmp/test_behavior_{i}.py") for i in range(100)}
+        assert tf._seen_paths == expected_paths