unit tests

Author: Codeflash Bot

codeflash/cli_cmds/cmd_init.py

@@ -53,7 +53,7 @@
 
 
 @dataclass(frozen=True)
-class SetupInfo:
+class CLISetupInfo:
     module_root: str
     tests_root: str
     benchmarks_root: Union[str, None]
@@ -101,7 +101,7 @@ def init_codeflash() -> None:
         git_remote = config.get("git_remote", "origin") if config else "origin"
 
         if should_modify:
-            setup_info: SetupInfo = collect_setup_info()
+            setup_info: CLISetupInfo = collect_setup_info()
             git_remote = setup_info.git_remote
             configured = configure_pyproject_toml(setup_info)
             if not configured:
@@ -240,7 +240,7 @@ def get_toml_key(self) -> str:
 
 
 @lru_cache(maxsize=1)
-def get_valid_subdirs() -> list[str]:
+def get_valid_subdirs(current_dir: Optional[Path] = None) -> list[str]:
     ignore_subdirs = [
         "venv",
         "node_modules",
@@ -253,8 +253,11 @@ def get_valid_subdirs() -> list[str]:
         "tmp",
         "__pycache__",
     ]
+    path_str = str(current_dir) if current_dir else "."
     return [
-        d for d in next(os.walk("."))[1] if not d.startswith(".") and not d.startswith("__") and d not in ignore_subdirs
+        d
+        for d in next(os.walk(path_str))[1]
+        if not d.startswith(".") and not d.startswith("__") and d not in ignore_subdirs
     ]
 
 
@@ -274,7 +277,7 @@ def get_suggestions(section: str) -> tuple(list[str], Optional[str]):
     raise ValueError(msg)
 
 
-def collect_setup_info() -> SetupInfo:
+def collect_setup_info() -> CLISetupInfo:
     curdir = Path.cwd()
     # Check if the cwd is writable
     if not os.access(curdir, os.W_OK):
@@ -554,7 +557,7 @@ def collect_setup_info() -> SetupInfo:
     enable_telemetry = ask_for_telemetry()
 
     ignore_paths: list[str] = []
-    return SetupInfo(
+    return CLISetupInfo(
         module_root=str(module_root),
         tests_root=str(tests_root),
         benchmarks_root=str(benchmarks_root) if benchmarks_root else None,
@@ -1003,7 +1006,9 @@ def get_formatter_cmds(formatter: str) -> list[str]:
 
 
 # Create or update the pyproject.toml file with the Codeflash dependency & configuration
-def configure_pyproject_toml(setup_info: Union[VsCodeSetupInfo, SetupInfo], config_file: Optional[Path] = None) -> bool:
+def configure_pyproject_toml(
+    setup_info: Union[VsCodeSetupInfo, CLISetupInfo], config_file: Optional[Path] = None
+) -> bool:
     for_vscode = isinstance(setup_info, VsCodeSetupInfo)
     toml_path = config_file or Path.cwd() / "pyproject.toml"
     try:

tests/test_cmd_init.py

@@ -0,0 +1,189 @@
+import pytest
+import tempfile
+from pathlib import Path
+from codeflash.cli_cmds.cmd_init import (
+    is_valid_pyproject_toml,
+    configure_pyproject_toml,
+    CLISetupInfo,
+    get_formatter_cmds,
+    VsCodeSetupInfo,
+    get_valid_subdirs,
+)
+import os
+
+
+@pytest.fixture
+def temp_dir():
+    with tempfile.TemporaryDirectory() as tmpdirname:
+        yield Path(tmpdirname)
+
+
+def test_is_valid_pyproject_toml_with_empty_config(temp_dir: Path) -> None:
+    with (temp_dir / "pyproject.toml").open(mode="w") as f:
+        f.write(
+            """[tool.codeflash]
+"""
+        )
+        f.flush()
+        valid, _, _message = is_valid_pyproject_toml(temp_dir / "pyproject.toml")
+        assert not valid
+        assert _message == "Missing required field: 'module_root'"
+
+def test_is_valid_pyproject_toml_with_incorrect_module_root(temp_dir: Path) -> None:
+    with (temp_dir / "pyproject.toml").open(mode="w") as f:
+        wrong_module_root = temp_dir / "invalid_directory"
+        f.write(
+            f"""[tool.codeflash]
+module-root = "invalid_directory"
+"""
+        )
+        f.flush()
+        valid, config, _message = is_valid_pyproject_toml(temp_dir / "pyproject.toml")
+        assert not valid
+        assert _message == f"Invalid 'module_root': directory does not exist at {wrong_module_root}"
+
+
+def test_is_valid_pyproject_toml_with_incorrect_tests_root(temp_dir: Path) -> None:
+    with (temp_dir / "pyproject.toml").open(mode="w") as f:
+        wrong_tests_root = temp_dir / "incorrect_tests_root"
+        f.write(
+            f"""[tool.codeflash]
+module-root = "."
+tests-root = "incorrect_tests_root"
+"""
+        )
+        f.flush()
+        valid, config, _message = is_valid_pyproject_toml(temp_dir / "pyproject.toml")
+        assert not valid
+        assert _message == f"Invalid 'tests_root': directory does not exist at {wrong_tests_root}"
+
+
+def test_is_valid_pyproject_toml_with_valid_config(temp_dir: Path) -> None:
+    with (temp_dir / "pyproject.toml").open(mode="w") as f:
+        os.makedirs(temp_dir / "tests")
+        f.write(
+            """[tool.codeflash]
+module-root = "."
+tests-root = "tests"
+test-framework = "pytest"
+"""
+        )
+        f.flush()
+        valid, config, _message = is_valid_pyproject_toml(temp_dir / "pyproject.toml")
+        assert valid
+
+def test_get_formatter_cmd(temp_dir: Path) -> None:
+    assert get_formatter_cmds("black") == ["black $file"]
+    assert get_formatter_cmds("ruff") == ["ruff check --exit-zero --fix $file", "ruff format $file"]
+    assert get_formatter_cmds("disabled") == ["disabled"]
+    assert get_formatter_cmds("don't use a formatter") == ["disabled"]
+
+def test_configure_pyproject_toml_for_cli(temp_dir: Path) -> None:
+
+    pyproject_path = temp_dir / "pyproject.toml"
+
+    with (pyproject_path).open(mode="w") as f:
+        f.write("")
+        f.flush()
+        os.mkdir(temp_dir / "tests")
+        config = CLISetupInfo(
+            module_root=".",
+            tests_root="tests",
+            benchmarks_root=None,
+            test_framework="pytest",
+            ignore_paths=[],
+            formatter="black",
+            git_remote="origin",
+            enable_telemetry=False,
+        )
+
+        success = configure_pyproject_toml(config, pyproject_path)
+        assert success
+
+        config_content = pyproject_path.read_text()
+        assert """[tool.codeflash]
+# All paths are relative to this pyproject.toml's directory.
+module-root = "."
+tests-root = "tests"
+test-framework = "pytest"
+ignore-paths = []
+disable-telemetry = true
+formatter-cmds = ["black $file"]
+""" == config_content
+        valid, _, _ = is_valid_pyproject_toml(pyproject_path)
+        assert valid
+
+def test_configure_pyproject_toml_for_vscode_with_empty_config(temp_dir: Path) -> None:
+
+    pyproject_path = temp_dir / "pyproject.toml"
+
+    with (pyproject_path).open(mode="w") as f:
+        f.write("")
+        f.flush()
+        os.mkdir(temp_dir / "tests")
+        config = VsCodeSetupInfo(
+            module_root=".",
+            tests_root="tests",
+            test_framework="pytest",
+            formatter="black",
+        )
+
+        success = configure_pyproject_toml(config, pyproject_path)
+        assert success
+
+        config_content = pyproject_path.read_text()
+        assert """[tool.codeflash]
+module-root = "."
+tests-root = "tests"
+test-framework = "pytest"
+formatter-cmds = ["black $file"]
+""" == config_content
+        valid, _, _ = is_valid_pyproject_toml(pyproject_path)
+        assert valid
+
+def test_configure_pyproject_toml_for_vscode_with_existing_config(temp_dir: Path) -> None:
+    pyproject_path = temp_dir / "pyproject.toml"
+    
+    with (pyproject_path).open(mode="w") as f:
+        f.write("""[tool.codeflash]
+module-root = "codeflash"
+tests-root = "tests"
+benchmarks-root = "tests/benchmarks"
+test-framework = "pytest"
+formatter-cmds = ["disabled"]
+""")
+        f.flush()
+        os.mkdir(temp_dir / "tests")
+        config = VsCodeSetupInfo(
+            module_root=".",
+            tests_root="tests",
+            test_framework="pytest",
+            formatter="disabled",
+        )
+
+        success = configure_pyproject_toml(config, pyproject_path)
+        assert success
+
+        config_content = pyproject_path.read_text()
+        # the benchmarks-root shouldn't get overwritten
+        assert """[tool.codeflash]
+module-root = "."
+tests-root = "tests"
+benchmarks-root = "tests/benchmarks"
+test-framework = "pytest"
+formatter-cmds = ["disabled"]
+""" == config_content
+        valid, _, _ = is_valid_pyproject_toml(pyproject_path)
+        assert valid
+
+def test_get_valid_subdirs(temp_dir: Path) -> None:
+    os.mkdir(temp_dir / "dir1")
+    os.mkdir(temp_dir / "dir2")
+    os.mkdir(temp_dir / "__pycache__")
+    os.mkdir(temp_dir / ".git")
+    os.mkdir(temp_dir / "tests")
+
+    dirs = get_valid_subdirs(temp_dir)
+    assert "tests" in dirs
+    assert "dir1" in dirs
+    assert "dir2" in dirs