iop/scripts/agent_benchmark/rubric_test.py

from __future__ import annotations

import hashlib
import json
import os
import tempfile
import unittest
from pathlib import Path

from scripts.agent_benchmark.manifest import (
    ONE_SHOT_RUBRIC_VERSION,
    RUBRIC_VERSION,
)
from scripts.agent_benchmark.rubric import (
    ONE_SHOT_RUBRIC_CATEGORIES,
    RUBRIC_CATEGORIES,
    RUBRIC_CATEGORIES_BY_VERSION,
    RubricError,
    canonical_worksheet_bytes,
    load_worksheet,
    rubric_categories,
    validate_worksheet,
)


def _worksheet(rubric_version: str = RUBRIC_VERSION) -> dict:
    categories = [
        {
            "id": ident,
            "max_score": maximum,
            "score": maximum,
            "evidence": f"Evidence for {ident}.",
        }
        for ident, maximum in rubric_categories(rubric_version)
    ]
    return {
        "rubric_version": rubric_version,
        "categories": categories,
        "total": 100,
    }


class RubricTest(unittest.TestCase):
    def test_exact_categories_and_total_are_accepted(self):
        worksheet = validate_worksheet(_worksheet())
        self.assertEqual(worksheet.total, 100)
        self.assertEqual(
            [(item.id, item.max_score) for item in worksheet.categories],
            list(RUBRIC_CATEGORIES),
        )
        self.assertEqual(
            validate_worksheet(json.loads(canonical_worksheet_bytes(worksheet))),
            worksheet,
        )
        self.assertEqual(
            hashlib.sha256(canonical_worksheet_bytes(worksheet)).hexdigest(),
            "233727170226ab30409657007751c13616a16c83ca75700e15a2c8dce56a96ee",
        )

    def test_one_shot_rubric_exact_categories_and_total_are_accepted(self):
        expected = (
            ("requirements_fidelity", 25),
            ("visual_completeness", 25),
            ("responsive_accessibility", 15),
            ("image_detail_usage", 10),
            ("behavior_stability", 10),
            ("code_quality", 10),
            ("self_verification", 5),
        )
        self.assertEqual(ONE_SHOT_RUBRIC_CATEGORIES, expected)
        self.assertEqual(rubric_categories(ONE_SHOT_RUBRIC_VERSION), expected)
        self.assertEqual(tuple(RUBRIC_CATEGORIES_BY_VERSION), (
            RUBRIC_VERSION,
            ONE_SHOT_RUBRIC_VERSION,
        ))
        self.assertEqual(sum(maximum for _, maximum in RUBRIC_CATEGORIES), 100)
        self.assertEqual(sum(maximum for _, maximum in expected), 100)

        worksheet = validate_worksheet(
            _worksheet(ONE_SHOT_RUBRIC_VERSION),
            expected_version=ONE_SHOT_RUBRIC_VERSION,
        )
        self.assertEqual(worksheet.rubric_version, ONE_SHOT_RUBRIC_VERSION)
        self.assertEqual(worksheet.total, 100)
        self.assertEqual(
            [(item.id, item.max_score) for item in worksheet.categories],
            list(expected),
        )

    def test_unknown_and_cross_version_worksheets_are_rejected(self):
        with self.assertRaises(RubricError):
            rubric_categories("unknown-rubric-v1")
        with self.assertRaises(RubricError):
            validate_worksheet(_worksheet(), expected_version=ONE_SHOT_RUBRIC_VERSION)
        with self.assertRaises(RubricError):
            validate_worksheet(
                _worksheet(ONE_SHOT_RUBRIC_VERSION),
                expected_version=RUBRIC_VERSION,
            )
        unknown = _worksheet()
        unknown["rubric_version"] = "unknown-rubric-v1"
        with self.assertRaises(RubricError):
            validate_worksheet(unknown)

    def test_cross_version_worksheet_is_rejected(self):
        with tempfile.TemporaryDirectory() as temp:
            path = Path(temp) / "worksheet.json"
            path.write_text(json.dumps(_worksheet()), encoding="utf-8")
            with self.assertRaises(RubricError):
                load_worksheet(path, expected_version=ONE_SHOT_RUBRIC_VERSION)

    def test_missing_extra_reordered_and_out_of_range_values_fail(self):
        cases = []
        missing = _worksheet()
        missing["categories"] = missing["categories"][:-1]
        cases.append(missing)
        extra = _worksheet()
        extra["automatic_gate_points"] = 1
        cases.append(extra)
        reordered = _worksheet()
        reordered["categories"] = list(reversed(reordered["categories"]))
        cases.append(reordered)
        too_high = _worksheet()
        too_high["categories"][0]["score"] = 26
        too_high["total"] = 101
        cases.append(too_high)
        bad_total = _worksheet()
        bad_total["total"] = 0
        cases.append(bad_total)
        empty_evidence = _worksheet()
        empty_evidence["categories"][0]["evidence"] = "  "
        cases.append(empty_evidence)
        for value in cases:
            with self.subTest(value=value):
                with self.assertRaises(RubricError):
                    validate_worksheet(value)

    def test_loader_refuses_nonregular_and_malformed_files(self):
        with tempfile.TemporaryDirectory() as temp:
            root = Path(temp)
            valid = root / "worksheet.json"
            valid.write_bytes(json.dumps(_worksheet()).encode())
            self.assertEqual(load_worksheet(valid).total, 100)

            malformed = root / "malformed.json"
            malformed.write_bytes(b"{}")
            with self.assertRaises(RubricError):
                load_worksheet(malformed)

            link = root / "link.json"
            os.symlink(valid, link)
            with self.assertRaises(RubricError):
                load_worksheet(link)


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