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.

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():

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)