484 lines
18 KiB
Python
484 lines
18 KiB
Python
"""
|
|
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.
|
|
|
|
The key difference between behavior and performance modes:
|
|
- behavior: Uses codeflash.capture() which writes to SQLite with full args/return values
|
|
- performance: Uses codeflash.capturePerf() which only prints timing to stdout
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import re
|
|
from pathlib import Path
|
|
from typing import Literal
|
|
|
|
|
|
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,
|
|
mode: Literal["behavior", "performance"] = "behavior",
|
|
) -> 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.
|
|
|
|
Static identifiers (funcName, lineId) are determined at instrumentation time.
|
|
The lineId enables tracking when the same call site is invoked multiple times (e.g., in loops).
|
|
|
|
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
|
|
mode: Testing mode - 'behavior' uses capture() for SQLite,
|
|
'performance' uses capturePerf() for stdout timing
|
|
|
|
Returns:
|
|
Instrumented test source code
|
|
"""
|
|
# Check if already instrumented
|
|
if "codeflash-jest-helper" in test_source:
|
|
# If already instrumented, convert between modes if needed
|
|
if mode == "performance" and "codeflash.capture(" in test_source and "codeflash.capturePerf(" not in test_source:
|
|
test_source = test_source.replace("codeflash.capture(", "codeflash.capturePerf(")
|
|
elif mode == "behavior" and "codeflash.capturePerf(" in test_source:
|
|
test_source = test_source.replace("codeflash.capturePerf(", "codeflash.capture(")
|
|
return test_source
|
|
|
|
lines = test_source.split("\n")
|
|
result_lines = []
|
|
|
|
# Add helper import at the top, after any existing imports
|
|
# Use relative path since helper is in project root and tests are in tests/ subdirectory
|
|
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)
|
|
|
|
# Choose the capture function based on mode
|
|
# - behavior: capture() writes to SQLite with args/return values
|
|
# - performance: capturePerf() prints timing markers to stdout (no SQLite overhead)
|
|
capture_func = "capturePerf" if mode == "performance" else "capture"
|
|
|
|
# Wrap function calls with codeflash.capture or codeflash.capturePerf
|
|
# 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*\(([^)]*)\)"
|
|
|
|
# Apply replacement line by line to track line numbers
|
|
# This enables tracking when the same call site is invoked multiple times (e.g., in loops)
|
|
instrumented_source = _safe_replace_function_calls_with_lineid(
|
|
instrumented_source, function_name, capture_func, pattern
|
|
)
|
|
|
|
# Comment out expect() assertions - we use captured behavior for verification instead
|
|
# This prevents hardcoded assertions from failing during baseline capture
|
|
instrumented_source = _comment_out_expects(instrumented_source)
|
|
|
|
return instrumented_source
|
|
|
|
|
|
def _comment_out_expects(source: str) -> str:
|
|
"""
|
|
Comment out expect() assertions in the test code.
|
|
|
|
During behavior capture, we want to execute the function calls but not
|
|
fail on hardcoded assertions. The captured outputs will be used for
|
|
verification instead.
|
|
|
|
IMPORTANT: If an expect() contains a codeflash.capture/capturePerf call,
|
|
we need to preserve that call while disabling the assertion.
|
|
e.g., expect(codeflash.capturePerf('fn', fn, arg)).toBe(x)
|
|
becomes: codeflash.capturePerf('fn', fn, arg); // [codeflash-disabled] .toBe(x)
|
|
"""
|
|
import re
|
|
|
|
# Pattern 1: expect() containing codeflash capture calls
|
|
# Transform: expect(codeflash.capture(...)).toBe/toEqual/etc(...)
|
|
# Into: codeflash.capture(...); // [codeflash-disabled] .toBe/toEqual/etc(...)
|
|
#
|
|
# Match: expect(codeflash.capture[Perf](...)).matcher(...)
|
|
# We need to handle nested parentheses in the capture call
|
|
lines = source.split('\n')
|
|
result_lines = []
|
|
|
|
for line in lines:
|
|
stripped = line.strip()
|
|
|
|
# Check if this line has expect(codeflash.capture... pattern
|
|
# Simple pattern: expect(codeflash.capture
|
|
if 'expect(codeflash.capture' in stripped:
|
|
# Extract the codeflash.capture(...) call and convert
|
|
# expect(codeflash.capturePerf('fn', fn, args)).toBe(expected);
|
|
# -> codeflash.capturePerf('fn', fn, args); // [codeflash-disabled] .toBe(expected);
|
|
|
|
# Find the capture call start
|
|
capture_match = re.search(r'expect\((codeflash\.capture(?:Perf)?)\(', line)
|
|
if capture_match:
|
|
indent = len(line) - len(line.lstrip())
|
|
capture_start = capture_match.end() - 1 # Position of opening paren of capture call
|
|
|
|
# Find matching closing paren for the capture call
|
|
paren_count = 1
|
|
i = capture_start + 1
|
|
while i < len(line) and paren_count > 0:
|
|
if line[i] == '(':
|
|
paren_count += 1
|
|
elif line[i] == ')':
|
|
paren_count -= 1
|
|
i += 1
|
|
|
|
# Now i points to just after the closing paren of capture call
|
|
capture_call = line[capture_match.start() + 7:i] # Skip "expect(" to get "codeflash.capture...)"
|
|
|
|
# Find what comes after the capture call
|
|
# rest_of_line is like ").toBe(expected);" - starts with ) from expect()
|
|
rest_of_line = line[i:]
|
|
|
|
# Check if rest starts with ).to or similar assertion
|
|
# The first ) is closing expect(), then we have .toBe(...)
|
|
assertion_match = re.match(r'\)(\.to\w+\([^)]*\);?)', rest_of_line)
|
|
if assertion_match:
|
|
# Extract just the assertion part without the leading )
|
|
assertion = assertion_match.group(1)
|
|
result_lines.append(' ' * indent + capture_call + '; // [codeflash-disabled] ' + assertion)
|
|
continue
|
|
elif rest_of_line.strip().startswith(')'):
|
|
# Just closing paren, no assertion - the capture call becomes standalone
|
|
result_lines.append(' ' * indent + capture_call + ';')
|
|
continue
|
|
|
|
# Pattern 2: Lines that start with expect( (but not codeflash calls)
|
|
if stripped.startswith('expect(') and 'codeflash.capture' not in stripped:
|
|
indent = len(line) - len(line.lstrip())
|
|
result_lines.append(' ' * indent + '// [codeflash-disabled] ' + stripped)
|
|
continue
|
|
|
|
# Pattern 3: expect() after semicolon on same line (but not codeflash calls)
|
|
if '; expect(' in line and 'codeflash.capture' not in line:
|
|
line = re.sub(r';\s*expect\(', '; // [codeflash-disabled] expect(', line)
|
|
|
|
result_lines.append(line)
|
|
|
|
return '\n'.join(result_lines)
|
|
|
|
|
|
def _safe_replace_function_calls_with_lineid(
|
|
source: str,
|
|
function_name: str,
|
|
capture_func: str,
|
|
pattern: str,
|
|
) -> str:
|
|
"""
|
|
Replace function calls while tracking line numbers.
|
|
|
|
Each function call is wrapped with a static lineId that enables tracking
|
|
when the same call site is invoked multiple times (e.g., in loops).
|
|
|
|
Args:
|
|
source: The JavaScript source code
|
|
function_name: Name of the function to wrap
|
|
capture_func: The capture function to use ('capture' or 'capturePerf')
|
|
pattern: Regex pattern to match function calls
|
|
"""
|
|
lines = source.split('\n')
|
|
result_lines = []
|
|
|
|
for line_num, line in enumerate(lines, start=1):
|
|
# Process each line independently to track line numbers
|
|
# Use the line number as the lineId
|
|
|
|
# Skip lines that are in strings or comments (simplified check)
|
|
if line.strip().startswith('//') or line.strip().startswith('/*'):
|
|
result_lines.append(line)
|
|
continue
|
|
|
|
# Track position within the line
|
|
processed_line = _replace_calls_in_line(
|
|
line, function_name, capture_func, pattern, line_num
|
|
)
|
|
result_lines.append(processed_line)
|
|
|
|
return '\n'.join(result_lines)
|
|
|
|
|
|
def _replace_calls_in_line(
|
|
line: str,
|
|
function_name: str,
|
|
capture_func: str,
|
|
pattern: str,
|
|
line_num: int,
|
|
) -> str:
|
|
"""
|
|
Replace function calls in a single line with capture wrapper.
|
|
|
|
Handles string literals and avoids replacing inside them.
|
|
"""
|
|
result = []
|
|
i = 0
|
|
length = len(line)
|
|
call_index = 0 # Track multiple calls on the same line
|
|
|
|
while i < length:
|
|
char = line[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 line[i] == "\\":
|
|
result.append(line[i:i+2] if i+1 < length else line[i:])
|
|
i += min(2, length - i)
|
|
elif line[i] == "$" and i + 1 < length and line[i + 1] == "{":
|
|
# Template expression - need to handle nested braces
|
|
result.append(line[i:i+2])
|
|
i += 2
|
|
brace_count = 1
|
|
while i < length and brace_count > 0:
|
|
if line[i] == "{":
|
|
brace_count += 1
|
|
elif line[i] == "}":
|
|
brace_count -= 1
|
|
result.append(line[i])
|
|
i += 1
|
|
elif line[i] == quote_char:
|
|
result.append(line[i])
|
|
i += 1
|
|
break
|
|
else:
|
|
result.append(line[i])
|
|
i += 1
|
|
else:
|
|
# Regular string
|
|
while i < length:
|
|
if line[i] == "\\":
|
|
result.append(line[i:i+2] if i+1 < length else line[i:])
|
|
i += min(2, length - i)
|
|
elif line[i] == quote_char:
|
|
result.append(line[i])
|
|
i += 1
|
|
break
|
|
else:
|
|
result.append(line[i])
|
|
i += 1
|
|
continue
|
|
|
|
# Check for function call pattern
|
|
remaining = line[i:]
|
|
match = re.match(pattern, remaining)
|
|
if match:
|
|
# Check that we're not preceded by a dot (method call) or already wrapped
|
|
if i > 0 and line[i - 1] == ".":
|
|
result.append(char)
|
|
i += 1
|
|
continue
|
|
|
|
# Check we haven't already wrapped this
|
|
check_start = max(0, i - 20)
|
|
preceding = line[check_start:i]
|
|
if "codeflash.capture" in preceding or "codeflash.capturePerf" in preceding:
|
|
result.append(char)
|
|
i += 1
|
|
continue
|
|
|
|
# Generate lineId: line_number_call_index
|
|
# For multiple calls on same line, we add call_index to distinguish them
|
|
line_id = f"{line_num}_{call_index}" if call_index > 0 else str(line_num)
|
|
call_index += 1
|
|
|
|
# Build replacement: codeflash.capture('funcName', 'lineId', func, args)
|
|
args = match.group(1).strip()
|
|
if args:
|
|
replacement = f"codeflash.{capture_func}('{function_name}', '{line_id}', {function_name}, {args})"
|
|
else:
|
|
replacement = f"codeflash.{capture_func}('{function_name}', '{line_id}', {function_name})"
|
|
|
|
result.append(replacement)
|
|
i += match.end()
|
|
continue
|
|
|
|
result.append(char)
|
|
i += 1
|
|
|
|
return "".join(result)
|
|
|
|
|
|
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}';
|
|
"""
|