From 26709f9b8266f60c046017393ceabdc4d5e9d6e7 Mon Sep 17 00:00:00 2001
From: Kevin Turcios <turcioskevinr@gmail.com>
Date: Tue, 21 Apr 2026 06:18:27 -0500
Subject: [PATCH] Fix test discovery for src-layout projects with
 module_root=src

The pytest collection subprocess only added module_root's parent to
PYTHONPATH, which works when module_root is a package (e.g. src/aviary)
but fails when it is the source root itself (e.g. src). Now both
module_root and its parent are added so imports like
`from mypkg.core import func` resolve correctly in either layout.
---
 .../test_discovery/discovery.py               | 15 ++++++--
 .../tests/test_test_discovery.py              | 38 +++++++++++++++++++
 2 files changed, 49 insertions(+), 4 deletions(-)

diff --git a/packages/codeflash-python/src/codeflash_python/test_discovery/discovery.py b/packages/codeflash-python/src/codeflash_python/test_discovery/discovery.py
index 8877bef..de6fb4c 100644
--- a/packages/codeflash-python/src/codeflash_python/test_discovery/discovery.py
+++ b/packages/codeflash-python/src/codeflash_python/test_discovery/discovery.py
@@ -138,15 +138,22 @@ def discover_tests_pytest(  # noqa: C901, PLR0912, PLR0915
         Path(__file__).parent.parent / "analysis" / "_discovery_worker.py"
     )
 
-    # For src-layout projects (e.g., module-root = src/aviary),
-    # add module_root's parent to PYTHONPATH so pytest can
-    # import the package (e.g., ``from aviary import ...``).
+    # Ensure PYTHONPATH covers both module_root and its parent so
+    # pytest can import packages in src-layout projects regardless
+    # of whether module_root is the source root (``src``) or a
+    # package inside it (``src/aviary``).
     env = os.environ.copy()
     if cfg.module_root is not None:
+        module_root_str = str(Path(cfg.module_root))
         parent = str(Path(cfg.module_root).parent)
+        extra = (
+            f"{module_root_str}{os.pathsep}{parent}"
+            if module_root_str != parent
+            else parent
+        )
         existing = env.get("PYTHONPATH", "")
         env["PYTHONPATH"] = (
-            f"{parent}{os.pathsep}{existing}" if existing else parent
+            f"{extra}{os.pathsep}{existing}" if existing else extra
         )
 
     with custom_addopts():
diff --git a/packages/codeflash-python/tests/test_test_discovery.py b/packages/codeflash-python/tests/test_test_discovery.py
index 08007cd..4455b18 100644
--- a/packages/codeflash-python/tests/test_test_discovery.py
+++ b/packages/codeflash-python/tests/test_test_discovery.py
@@ -2256,3 +2256,41 @@ def test_discover_unit_tests_caching():
     assert (
         non_cached_num_discovered_replay_tests == num_discovered_replay_tests
     )
+
+
+def test_discover_tests_src_layout_module_root():
+    """Tests are discovered when module_root is the src directory itself."""
+    with tempfile.TemporaryDirectory() as tmp:
+        root = Path(tmp)
+        src = root / "src"
+        pkg = src / "mypkg"
+        pkg.mkdir(parents=True)
+        (pkg / "__init__.py").write_text("")
+        (pkg / "core.py").write_text("def do_work():\n    return 42\n")
+
+        tests_dir = root / "tests"
+        tests_dir.mkdir()
+        (tests_dir / "test_core.py").write_text(
+            "from mypkg.core import do_work\n\n"
+            "def test_do_work():\n"
+            "    assert 42 == do_work()\n"
+        )
+
+        func = FunctionToOptimize(
+            function_name="do_work",
+            file_path=pkg / "core.py",
+            parents=(),
+        )
+        cfg = TestConfig(
+            tests_project_rootdir=root,
+            tests_root=str(tests_dir),
+            project_root_path=str(root),
+            test_framework="pytest",
+            module_root=src,
+        )
+        result, n_tests, _n_replay = discover_unit_tests(
+            cfg,
+            file_to_funcs_to_optimize={func.file_path: [func]},
+        )
+        assert 1 <= n_tests
+        assert any("do_work" in key for key in result)