codeflash/tests/benchmarks/test_benchmark_code_extract_code_context.py
Kevin Turcios 547c02e8bc refactor: move context extraction modules to languages/python/context/
Move code_context_extractor.py and unused_definition_remover.py from
codeflash/context/ to codeflash/languages/python/context/ and update
all import sites.
2026-02-16 14:49:04 -05:00

31 lines
1.2 KiB
Python

from argparse import Namespace
from pathlib import Path
from codeflash.discovery.functions_to_optimize import FunctionToOptimize
from codeflash.languages.python.context.code_context_extractor import get_code_optimization_context
from codeflash.models.models import FunctionParent
from codeflash.optimization.optimizer import Optimizer
def test_benchmark_extract(benchmark) -> None:
file_path = Path(__file__).parent.parent.parent.resolve() / "codeflash"
opt = Optimizer(
Namespace(
project_root=file_path.resolve(),
disable_telemetry=True,
tests_root=(file_path / "tests").resolve(),
test_framework="pytest",
pytest_cmd="pytest",
experiment_id=None,
test_project_root=Path.cwd(),
)
)
function_to_optimize = FunctionToOptimize(
function_name="replace_function_and_helpers_with_optimized_code",
file_path=file_path / "optimization" / "function_optimizer.py",
parents=[FunctionParent(name="FunctionOptimizer", type="ClassDef")],
starting_line=None,
ending_line=None,
)
benchmark(get_code_optimization_context, function_to_optimize, opt.args.project_root)