codeflash-internal/django/aiservice/testgen/instrumentation/javascript/instrument_javascript.py

"""
JavaScript test instrumentation module.

This module instruments JavaScript tests by injecting the codeflash-jest-helper
to capture function call behavior and performance data.

Unlike Python which has separate instrumentation for generated vs existing tests,
JavaScript uses a UNIFIED approach - the same instrumentation works for all tests.
"""

from __future__ import annotations

import re
from pathlib import Path


def get_jest_helper_path() -> Path:
    """
    Get the path to the codeflash-jest-helper.js file.

    Returns:
        Path to the helper JavaScript file
    """
    return Path(__file__).parent / "codeflash-jest-helper.js"


def instrument_javascript_tests(
    test_source: str,
    function_name: str,
    module_path: str,
) -> str:
    """
    Instrument JavaScript tests with codeflash helper.

    This is a UNIFIED approach - works for both generated and existing tests.
    The instrumentation wraps function calls to capture inputs, outputs, and timing.

    Args:
        test_source: The JavaScript test source code
        function_name: The name of the function being tested
        module_path: The path to the module containing the function

    Returns:
        Instrumented test source code
    """
    # Check if already instrumented
    if "codeflash-jest-helper" in test_source:
        return test_source

    lines = test_source.split("\n")
    result_lines = []

    # Add helper import at the top, after any existing imports
    helper_import = "const codeflash = require('codeflash-jest-helper');"
    import_inserted = False
    in_import_block = False

    for i, line in enumerate(lines):
        stripped = line.strip()

        # Track if we're in the import block
        if stripped.startswith("import ") or stripped.startswith("const ") and "require" in stripped:
            in_import_block = True
        elif in_import_block and stripped and not stripped.startswith("import ") and not (
            stripped.startswith("const ") and "require" in stripped
        ):
            # End of import block - insert helper import
            if not import_inserted:
                result_lines.append(helper_import)
                result_lines.append("")
                import_inserted = True
            in_import_block = False

        result_lines.append(line)

    # If no imports found, add at the beginning
    if not import_inserted:
        result_lines = [helper_import, ""] + result_lines

    instrumented_source = "\n".join(result_lines)

    # Wrap function calls with codeflash.capture
    # Pattern matches: functionName(args) but not inside strings or comments
    # This is a simple approach - a more robust solution would use an AST parser

    # Pattern for standalone function calls (not method calls)
    pattern = rf"(?<![.\w]){re.escape(function_name)}\s*\(([^)]*)\)"

    def replace_call(match: re.Match) -> str:
        args = match.group(1).strip()
        if args:
            return f"codeflash.capture('{function_name}', {function_name}, {args})"
        else:
            return f"codeflash.capture('{function_name}', {function_name})"

    # Apply replacement carefully - avoid replacing inside strings
    # This is a simplified approach that works for most cases
    instrumented_source = _safe_replace_function_calls(
        instrumented_source, function_name, replace_call, pattern
    )

    return instrumented_source


def _safe_replace_function_calls(
    source: str,
    function_name: str,
    replace_func: callable,
    pattern: str,
) -> str:
    """
    Replace function calls while avoiding string literals and comments.

    This is a simplified approach that handles common cases.
    A more robust solution would use a proper JavaScript parser.
    """
    result = []
    i = 0
    length = len(source)

    while i < length:
        char = source[i]

        # Skip string literals
        if char in "'\"`":
            quote_char = char
            result.append(char)
            i += 1

            # Handle template literals with ${} expressions
            if quote_char == "`":
                while i < length:
                    if source[i] == "\\":
                        result.append(source[i:i+2])
                        i += 2
                    elif source[i] == "$" and i + 1 < length and source[i + 1] == "{":
                        # Template expression - need to handle nested braces
                        result.append(source[i:i+2])
                        i += 2
                        brace_count = 1
                        while i < length and brace_count > 0:
                            if source[i] == "{":
                                brace_count += 1
                            elif source[i] == "}":
                                brace_count -= 1
                            result.append(source[i])
                            i += 1
                    elif source[i] == quote_char:
                        result.append(source[i])
                        i += 1
                        break
                    else:
                        result.append(source[i])
                        i += 1
            else:
                # Regular string
                while i < length:
                    if source[i] == "\\":
                        result.append(source[i:i+2])
                        i += 2
                    elif source[i] == quote_char:
                        result.append(source[i])
                        i += 1
                        break
                    else:
                        result.append(source[i])
                        i += 1
            continue

        # Skip single-line comments
        if char == "/" and i + 1 < length and source[i + 1] == "/":
            while i < length and source[i] != "\n":
                result.append(source[i])
                i += 1
            continue

        # Skip multi-line comments
        if char == "/" and i + 1 < length and source[i + 1] == "*":
            result.append(source[i:i+2])
            i += 2
            while i < length - 1:
                if source[i] == "*" and source[i + 1] == "/":
                    result.append(source[i:i+2])
                    i += 2
                    break
                result.append(source[i])
                i += 1
            continue

        # Check for function call pattern
        remaining = source[i:]
        match = re.match(pattern, remaining)
        if match:
            # Check that we're not preceded by a dot (method call)
            if i > 0 and source[i - 1] == ".":
                result.append(char)
                i += 1
                continue

            # Check that we haven't already wrapped this
            if i >= len("codeflash.capture") and source[i - len("codeflash.capture"):i] == "codeflash.capture":
                result.append(char)
                i += 1
                continue

            # Apply replacement
            replacement = replace_func(match)
            result.append(replacement)
            i += match.end()
            continue

        result.append(char)
        i += 1

    return "".join(result)


def get_jest_setup_code(output_file: str, mode: str = "behavior", loop_index: int = 0) -> str:
    """
    Generate Jest setup code for setting environment variables.

    Args:
        output_file: Path where results should be written
        mode: Testing mode ('behavior' or 'performance')
        loop_index: Current benchmark loop iteration

    Returns:
        JavaScript code to set up the testing environment
    """
    return f"""
// Codeflash test setup
process.env.CODEFLASH_OUTPUT_FILE = '{output_file}';
process.env.CODEFLASH_MODE = '{mode}';
process.env.CODEFLASH_LOOP_INDEX = '{loop_index}';
"""