concolic_testing.py

from __future__ import annotations

import ast
import os
import subprocess
import tempfile
import time
from pathlib import Path
from typing import TYPE_CHECKING

from codeflash.cli_cmds.console import console, logger
from codeflash.code_utils.compat import SAFE_SYS_EXECUTABLE
from codeflash.code_utils.concolic_utils import clean_concolic_tests, is_valid_concolic_test
from codeflash.code_utils.static_analysis import has_typed_parameters
from codeflash.discovery.discover_unit_tests import discover_unit_tests
from codeflash.languages import is_python
from codeflash.lsp.helpers import is_LSP_enabled
from codeflash.telemetry.posthog_cf import ph
from codeflash.verification.verification_utils import TestConfig

if TYPE_CHECKING:
    from argparse import Namespace

    from codeflash.discovery.functions_to_optimize import FunctionToOptimize
    from codeflash.models.models import FunctionCalledInTest


def generate_concolic_tests(
    test_cfg: TestConfig, args: Namespace, function_to_optimize: FunctionToOptimize, function_to_optimize_ast: ast.AST
) -> tuple[dict[str, set[FunctionCalledInTest]], str]:
    """Generate concolic tests using CrossHair (Python only).

    CrossHair is a Python-specific symbolic execution tool. For non-Python languages
    (JavaScript, TypeScript, etc.), this function returns early with empty results.

    Args:
        test_cfg: Test configuration
        args: Command line arguments
        function_to_optimize: The function being optimized
        function_to_optimize_ast: AST of the function (Python ast.FunctionDef)

    Returns:
        Tuple of (function_to_tests mapping, concolic test suite code)

    """
    start_time = time.perf_counter()
    function_to_concolic_tests = {}
    concolic_test_suite_code = ""

    # CrossHair is Python-only - skip for other languages
    if not is_python():
        logger.debug("Skipping concolic test generation for non-Python languages (CrossHair is Python-only)")
        return function_to_concolic_tests, concolic_test_suite_code

    if is_LSP_enabled():
        logger.debug("Skipping concolic test generation in LSP mode")
        return function_to_concolic_tests, concolic_test_suite_code

    if (
        test_cfg.concolic_test_root_dir
        and isinstance(function_to_optimize_ast, ast.FunctionDef)
        and has_typed_parameters(function_to_optimize_ast, function_to_optimize.parents)
    ):
        logger.info("Generating concolic opcode coverage tests for the original code…")
        console.rule()
        try:
            env = os.environ.copy()
            pythonpath = env.get("PYTHONPATH", "")
            project_root_str = str(args.project_root)
            if pythonpath:
                env["PYTHONPATH"] = f"{project_root_str}{os.pathsep}{pythonpath}"
            else:
                env["PYTHONPATH"] = project_root_str
            cover_result = subprocess.run(
                [
                    SAFE_SYS_EXECUTABLE,
                    "-m",
                    "crosshair",
                    "cover",
                    "--example_output_format=pytest",
                    "--per_condition_timeout=20",
                    ".".join(
                        [
                            function_to_optimize.file_path.relative_to(args.project_root)
                            .with_suffix("")
                            .as_posix()
                            .replace("/", "."),
                            function_to_optimize.qualified_name,
                        ]
                    ),
                ],
                capture_output=True,
                text=True,
                cwd=args.project_root,
                check=False,
                timeout=600,
                env=env,
            )
        except subprocess.TimeoutExpired:
            logger.debug("CrossHair Cover test generation timed out")
            return function_to_concolic_tests, concolic_test_suite_code

        if cover_result.returncode == 0:
            generated_concolic_test: str = cover_result.stdout
            if not is_valid_concolic_test(generated_concolic_test, project_root=str(args.project_root)):
                logger.debug("CrossHair generated invalid test, skipping")
                console.rule()
                return function_to_concolic_tests, concolic_test_suite_code
            concolic_test_suite_code: str = clean_concolic_tests(generated_concolic_test)
            concolic_test_suite_dir = Path(tempfile.mkdtemp(dir=test_cfg.concolic_test_root_dir))
            concolic_test_suite_path = concolic_test_suite_dir / "test_concolic_coverage.py"
            concolic_test_suite_path.write_text(concolic_test_suite_code, encoding="utf8")

            concolic_test_cfg = TestConfig(
                tests_root=concolic_test_suite_dir,
                tests_project_rootdir=test_cfg.concolic_test_root_dir,
                project_root_path=args.project_root,
            )
            function_to_concolic_tests, num_discovered_concolic_tests, _ = discover_unit_tests(concolic_test_cfg)
            logger.info(
                f"Created {num_discovered_concolic_tests} "
                f"concolic unit test case{'s' if num_discovered_concolic_tests != 1 else ''} "
            )
            console.rule()
            ph("cli-optimize-concolic-tests", {"num_tests": num_discovered_concolic_tests})

        else:
            logger.debug(f"Error running CrossHair Cover {': ' + cover_result.stderr if cover_result.stderr else '.'}")
            console.rule()
    end_time = time.perf_counter()
    logger.debug(f"Generated concolic tests in {end_time - start_time:.2f} seconds")
    return function_to_concolic_tests, concolic_test_suite_code