code/self_healing_finetuning/tests/test_repair_acceptance_gate.py

from __future__ import annotations

import json
import tempfile
import unittest
from pathlib import Path

from data_pipeline.repair_acceptance_gate import run_repair_acceptance_gate


def write_jsonl(path: Path, rows: list[dict[str, object]]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text("\n".join(json.dumps(row) for row in rows) + "\n", encoding="utf-8")


class RepairAcceptanceGateTests(unittest.TestCase):
    def test_accepts_clean_repairs_and_rejects_think_leakage(self) -> None:
        with tempfile.TemporaryDirectory() as tmp:
            root = Path(tmp)
            input_path = root / "repair_targets.jsonl"
            write_jsonl(
                input_path,
                [
                    {
                        "source_prediction_id": "good",
                        "failure_bucket": "accounting_sec_extraction",
                        "task": "finance_qa",
                        "prompt": "Revenue grew 4%, but backlog declined 8%. Summarize the revenue risk.",
                        "repair_answer": (
                            "Reported facts: revenue grew 4% while backlog declined 8%. "
                            "Inference: the backlog decline suggests future revenue pressure. "
                            "Risk/tradeoff: current growth may not persist if enterprise demand remains slower. "
                            "Conclusion: no direct risk ratio should be calculated from 4% and 8%."
                        ),
                    },
                    {
                        "source_prediction_id": "bad",
                        "failure_bucket": "accounting_sec_extraction",
                        "task": "finance_qa",
                        "prompt": "Revenue grew 4%, but backlog declined 8%. Summarize the revenue risk.",
                        "repair_answer": "<think>Okay, let's see. The user wants...</think> This is bullish.",
                    },
                ],
            )
            result = run_repair_acceptance_gate(input_path=input_path, output_dir=root / "gate", min_acceptance_rate=0.5)
            self.assertTrue(result["ok"], result)
            self.assertEqual(result["accepted_count"], 1)
            self.assertEqual(result["rejected_count"], 1)
            self.assertIn("no_think_tags", result["rejection_reason_counts"])
            self.assertTrue(Path(result["accepted_path"]).exists())
            self.assertTrue(Path(result["rejected_path"]).exists())


if __name__ == "__main__":
    unittest.main()