tests/test_study2_experiment.py

from __future__ import annotations

from pathlib import Path
import unittest

from agent_harness.specs import load_models
from agent_harness.study2_experiment import (
    Study2ExperimentError,
    _RuntimeLease,
    _json_object,
    _select_files,
    tokenizer_for,
)


ROOT = Path(__file__).resolve().parents[1]


class Study2ExperimentTests(unittest.TestCase):
    def test_agentless_json_and_file_selection_are_bounded(self) -> None:
        content = '{"files": ["a.go", "missing.go", "b.go"]}'
        self.assertEqual(_json_object(content)["files"][0], "a.go")
        self.assertEqual(_select_files(content, ("a.go", "b.go")), ("a.go", "b.go"))

    def test_each_study2_model_has_a_pinned_tokenizer(self) -> None:
        models = load_models(ROOT)
        for model_id in ("M002", "M003"):
            tokenizer = tokenizer_for(models[model_id])
            self.assertTrue(tokenizer.path.exists())
            self.assertGreater(tokenizer.count("repository navigation"), 0)

    def test_runtime_cleanup_does_not_mask_an_existing_exception(self) -> None:
        class BrokenResidency:
            def unload_all(self):
                raise RuntimeError("server connection lost")

        class StoppedServer:
            def status(self):
                return {"running": False, "returncode": 1}

        lease = _RuntimeLease(StoppedServer(), BrokenResidency(), True)  # type: ignore[arg-type]
        lease.__exit__(RuntimeError, RuntimeError("original"), None)
        self.assertEqual(lease.stop_state["action"], "already_stopped")
        self.assertIn("server connection lost", lease.cleanup_errors[0])

    def test_runtime_cleanup_failure_is_fatal_without_prior_exception(self) -> None:
        class BrokenResidency:
            def unload_all(self):
                raise RuntimeError("cannot unload")

        class StoppedServer:
            def status(self):
                return {"running": False, "returncode": 1}

        lease = _RuntimeLease(StoppedServer(), BrokenResidency(), True)  # type: ignore[arg-type]
        with self.assertRaises(Study2ExperimentError):
            lease.__exit__(None, None, None)


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