future house fixes

Author: KRRT7

tests/scripts/end_to_end_test_futurehouse.py

@@ -1,13 +1,18 @@
+from pathlib import Path
 import os
-import pathlib
 
-from end_to_end_test_utilities import CoverageExpectation, TestConfig, run_codeflash_command, run_with_retries
+from tests.scripts.end_to_end_test_utilities import (
+    CoverageExpectation,
+    TestConfig,
+    run_codeflash_command,
+    run_with_retries
+)
 
 
 def run_test(expected_improvement_pct: int) -> bool:
     config = TestConfig(
-        file_path="src/aviary/common_tags.py",
-        expected_unit_tests=0, # todo: fix bug https://linear.app/codeflash-ai/issue/CF-921/test-discovery-does-not-work-properly-for-e2e-futurehouse-example for context
+        file_path=Path("src/aviary/common_tags.py"),
+        expected_unit_tests=2,
         min_improvement_x=0.05,
         coverage_expectations=[
             CoverageExpectation(
@@ -17,10 +22,10 @@ def run_test(expected_improvement_pct: int) -> bool:
             )
         ],
     )
-    cwd = (
-        pathlib.Path(__file__).parent.parent.parent / "code_to_optimize" / "code_directories" / "futurehouse_structure"
+    future_house_working_dir = (
+        Path(__file__).parent.parent.parent / "code_to_optimize" / "code_directories" / "futurehouse_structure"
     ).resolve()
-    return run_codeflash_command(cwd, config, expected_improvement_pct)
+    return run_codeflash_command(future_house_working_dir, config, expected_improvement_pct)
 
 
 if __name__ == "__main__":

tests/scripts/end_to_end_test_utilities.py

@@ -7,6 +7,8 @@
 import time
 from dataclasses import dataclass, field
 from typing import Optional
+import contextlib
+import tomllib
 
 
 @dataclass
@@ -82,7 +84,6 @@ def run_codeflash_command(
     logging.basicConfig(level=logging.INFO)
     if config.trace_mode:
         return run_trace_test(cwd, config, expected_improvement_pct)
-
     path_to_file = cwd / config.file_path
     file_contents = path_to_file.read_text("utf-8")
     pytest_dir = cwd / "tests" / "pytest"
@@ -102,7 +103,6 @@ def run_codeflash_command(
 
     return_code = process.wait()
     stdout = "".join(output)
-
     validated = validate_output(stdout, return_code, expected_improvement_pct, config)
     if not validated:
         # Write original file contents back to file
@@ -129,7 +129,20 @@ def build_command(
 
     if config.function_name:
         base_command.extend(["--function", config.function_name])
-    base_command.extend(["--tests-root", str(test_root), "--module-root", str(cwd)])
+    
+    # Check if pyproject.toml exists with codeflash config - if so, don't override it
+    pyproject_path = cwd / "pyproject.toml"
+    has_codeflash_config = False
+    if pyproject_path.exists():
+        with contextlib.suppress(Exception):
+            with open(pyproject_path, "rb") as f:
+                pyproject_data = tomllib.load(f)
+                has_codeflash_config = "tool" in pyproject_data and "codeflash" in pyproject_data["tool"]
+    
+    # Only pass --tests-root and --module-root if they're not configured in pyproject.toml
+    if not has_codeflash_config:
+        base_command.extend(["--tests-root", str(test_root), "--module-root", str(cwd)])
+    
     if benchmarks_root:
         base_command.extend(["--benchmark", "--benchmarks-root", str(benchmarks_root)])
     if config.use_worktree: