CmdForge/tests/test_prompt_optimizer.py

"""Tests for prompt optimization (M9.1)."""

from unittest.mock import patch

from cmdforge.prompt_optimizer import (
    PromptVariation,
    OptimizationResult,
    generate_variations,
    optimize_tool,
    _get_strategies,
    _truncate,
)


class TestTruncate:
    def test_short_text_unchanged(self):
        assert _truncate("hello", 10) == "hello"

    def test_long_text_truncated(self):
        result = _truncate("this is a very long sentence that needs truncation", 20)
        assert len(result) <= 20
        assert result.endswith("...")


class TestGenerateVariations:
    def test_mock_provider_never_calls_external_provider(self):
        with patch("cmdforge.providers.call_provider") as call:
            _get_strategies("Summarize the text", 2, provider="mock")
        call.assert_not_called()

    def test_selected_provider_is_honored(self):
        from cmdforge.providers import ProviderResult
        with patch(
            "cmdforge.providers.call_provider",
            return_value=ProviderResult(text="A materially clearer instruction", success=True),
        ) as call:
            result = _get_strategies("Summarize the text", 1, provider="chosen")
        assert result[0][1] == "A materially clearer instruction"
        assert call.call_args.args[0] == "chosen"

    def test_generates_for_prompt_steps(self, tmp_path):
        from cmdforge.tool import Tool, PromptStep

        with patch("cmdforge.tool.TOOLS_DIR", tmp_path / ".cmdforge"):
            tool = Tool(
                name="test-tool",
                steps=[
                    PromptStep(
                        prompt="Analyze the input and produce a summary",
                        provider="mock",
                        output_var="summary",
                    ),
                ],
                output="{summary}",
            )
            variations = generate_variations(tool, count=3)
            assert len(variations) >= 1
            for v in variations:
                assert v.step_index == 0
                assert v.original == "Analyze the input and produce a summary"
                assert v.variation != v.original

    def test_skips_non_prompt_steps(self, tmp_path):
        from cmdforge.tool import Tool, CodeStep

        with patch("cmdforge.tool.TOOLS_DIR", tmp_path / ".cmdforge"):
            tool = Tool(
                name="test-tool",
                steps=[CodeStep(code="print('hello')", output_var="out")],
                output="{out}",
            )
            variations = generate_variations(tool)
            assert len(variations) == 0


class TestGetStrategies:
    def test_produces_variations(self):
        strategies = _get_strategies("Summarize the text", 3)
        assert len(strategies) == 3
        for strat, text in strategies:
            assert len(text) > 0
            assert strat in ("focused", "concise", "detailed", "rephrase", "structure")

    def test_truncates_concise(self):
        original = "A" * 200
        strategies = _get_strategies(original, 3)
        concise = [t for s, t in strategies if s == "concise"]
        if concise:
            assert len(concise[0]) <= 125  # 120 + "..."


class TestOptimizeTool:
    def test_structural_results_do_not_claim_a_best_prompt(self):
        from cmdforge.tool import Tool, PromptStep
        tool = Tool(
            name="structural", steps=[PromptStep("Summarize", "mock", "out")],
            output="{out}", input_schema={"type": "string"},
            output_schema={"type": "string"},
        )
        result = optimize_tool(tool, count=2)
        assert result.best is None
        assert "cannot compare prompt semantics" in result.note

    def test_behavioral_evaluator_selects_only_an_improvement(self):
        from cmdforge.tool import Tool, PromptStep
        tool = Tool(
            name="behavioral", steps=[PromptStep("Do it", "mock", "out")],
            output="{out}",
        )
        result = optimize_tool(
            tool, count=3,
            evaluator=lambda candidate: len(candidate.steps[0].prompt),
        )
        assert result.best is not None
        assert result.scores[result.best_index] > result.baseline_score
    def test_baseline_score_recorded(self, tmp_path):
        from cmdforge.tool import Tool, PromptStep

        with patch("cmdforge.tool.TOOLS_DIR", tmp_path / ".cmdforge"):
            tool = Tool(
                name="opt-test",
                steps=[
                    PromptStep(
                        prompt="Summarize: {input}",
                        provider="mock",
                        output_var="result",
                    ),
                ],
                output="{result}",
            )
            result = optimize_tool(tool, count=2)
            assert result.tool_name == "opt-test"
            assert result.baseline_score >= 0
            assert len(result.variations) >= 1
            assert len(result.scores) == len(result.variations)

    def test_no_prompt_steps_returns_empty(self, tmp_path):
        from cmdforge.tool import Tool, CodeStep

        with patch("cmdforge.tool.TOOLS_DIR", tmp_path / ".cmdforge"):
            tool = Tool(
                name="code-only",
                steps=[CodeStep(code="pass", output_var="x")],
                output="{x}",
            )
            result = optimize_tool(tool)
            assert len(result.variations) == 0
            assert result.best is None


class TestPromptVariation:
    def test_to_dict(self):
        v = PromptVariation(
            step_index=0,
            original="hello",
            variation="hi there",
            strategy="rephrase",
        )
        d = v.to_dict()
        assert d["original"] == "hello"
        assert d["strategy"] == "rephrase"


class TestOptimizationResult:
    def test_improvement_calculation(self):
        result = OptimizationResult(
            tool_name="test",
            variations=[
                PromptVariation(0, "a", "b", "rephrase"),
            ],
            scores={0: 5},
            best_index=0,
            baseline_score=3,
        )
        assert result.best is not None
        assert result.improvement > 0

    def test_no_improvement_when_no_baseline(self):
        result = OptimizationResult(
            tool_name="test",
            variations=[],
            scores={},
            baseline_score=0,
        )
        assert result.best is None
        assert result.improvement == 0.0