agent.py

"""Agent orchestration for TestForge."""

from __future__ import annotations

import shutil
import tempfile
from dataclasses import dataclass
from pathlib import Path

from analyzer import Analysis, analyze
from generator import GeneratedSuite, deterministic_inputs, generate_suite, write_suite
from inject import reset_sample
from model_suggest import suggest_extra_inputs
from mutator import MUTATIONS, Mutation, MutationScore, mutation_score
from patch import write_patch
from runner import RunResult, run_pytest


ROOT = Path(__file__).resolve().parent
SAMPLES_ROOT = ROOT / "samples"
LEGACY_ROOT = SAMPLES_ROOT / "legacy_repo"
RUNS_ROOT = ROOT / "generated" / "runs"


@dataclass(frozen=True)
class ForgeArtifacts:
    run_dir: Path
    analysis: Analysis
    suite: GeneratedSuite
    green: RunResult
    mutation: MutationScore
    patch_path: Path
    logs: tuple[str, ...]


def forge_legacy_repo(
    max_cases_per_function: int = 4,
    mutations: tuple[Mutation, ...] = MUTATIONS,
    use_model: bool = False,
) -> ForgeArtifacts:
    """Run the TestForge pipeline for the bundled sample.

    The deterministic path (`use_model=False`) is the demo backbone and needs
    no model or GPU. When `use_model=True`, a small coder model proposes extra
    edge-case argument tuples per function via `model_suggest.suggest_extra_inputs`;
    every suggestion is still captured by real execution before it becomes a test.
    """
    run_dir = _fresh_run_dir()
    logs: list[str] = []

    logs.append("Target: bundled legacy_repo with zero tests")
    reset_sample(SAMPLES_ROOT, run_dir)

    analysis = analyze(run_dir / "samples" / "legacy_repo", package="legacy_repo")
    logs.append(f"Analyzed {len(analysis.functions)} public functions")
    for fn in analysis.functions:
        logs.append(f"  - {fn.module}.{fn.qualname}({', '.join(fn.parameters)})")

    extra_inputs: dict[str, list[tuple]] = {}
    if use_model:
        logs.append("Model assist: asking small coder model for extra edge cases")
        for fn in analysis.functions:
            if fn.arity == 0:
                continue
            deterministic = deterministic_inputs(fn, max_cases=max_cases_per_function)
            suggestions = suggest_extra_inputs(fn, deterministic, max_new=2)
            if suggestions:
                extra_inputs[f"{fn.module}.{fn.qualname}"] = suggestions
                logs.append(f"  +{len(suggestions)} model-proposed case(s) for {fn.qualname}")
        if not extra_inputs:
            logs.append("  no extra cases proposed (model/GPU unavailable or nothing new found)")

    suite = generate_suite(
        analysis,
        max_cases_per_function=max_cases_per_function,
        extra_inputs=extra_inputs,
    )
    write_suite(suite, run_dir)
    logs.append(f"Captured {suite.assertion_count} behaviors into {len(suite.files)} test files")

    green = run_pytest(run_dir, package_parent=run_dir / "samples")
    logs.append(f"Suite run: {green.summary}")
    if green.coverage is not None:
        logs.append(f"Coverage: 0% -> {green.coverage:.0f}%")
    else:
        logs.append("Coverage: pytest-cov unavailable locally; suite still ran green")

    mutation = mutation_score(run_dir / "samples", suite, run_dir / "mutants", mutations=mutations)
    logs.append(f"Mutation score: {mutation.headline()} ({mutation.percent:.1f}%)")
    for survived in mutation.survived:
        logs.append(f"  survived: {survived.label}")

    patch_path = write_patch(suite, run_dir / "testforge.patch")
    logs.append(f"Patch ready: {patch_path.name}")

    return ForgeArtifacts(
        run_dir=run_dir,
        analysis=analysis,
        suite=suite,
        green=green,
        mutation=mutation,
        patch_path=patch_path,
        logs=tuple(logs),
    )


def _fresh_run_dir() -> Path:
    RUNS_ROOT.mkdir(parents=True, exist_ok=True)
    path = Path(tempfile.mkdtemp(prefix="testforge_", dir=RUNS_ROOT))
    if path.exists():
        shutil.rmtree(path)
    path.mkdir(parents=True)
    return path