tests/test_analysis.py

from __future__ import annotations

import unittest

from agent_harness.analysis import summarize_rows


class AnalysisTests(unittest.TestCase):
    def test_paired_treatment_summary(self) -> None:
        rows = []
        values = {
            "H000": [(0.0, 0.5), (1.0, 0.5)],
            "H001": [(1.0, 1.0), (1.0, 1.0)],
        }
        for harness, task_values in values.items():
            for index, (recall, mrr) in enumerate(task_values, start=1):
                rows.append(
                    {
                        "task_id": f"TASK_{index}",
                        "harness_id": harness,
                        "all_gold_in_top_10": recall == 1.0,
                        "file_recall_at_1": recall,
                        "file_recall_at_5": recall,
                        "file_recall_at_10": recall,
                        "mrr": mrr,
                        "ndcg_at_10": recall,
                        "query_seconds": 0.1,
                    }
                )
        summary = summarize_rows(rows)
        self.assertEqual(summary["treatments"]["H001"]["all_gold_in_top_10_count"], 2)
        self.assertEqual(
            summary["paired_differences"]["H001_minus_H000"]["mean_delta_file_recall_at_10"],
            0.5,
        )


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