submmit/tests/test_agent_entrypoint.py

from __future__ import annotations

import importlib.util
import os
import unittest
import tempfile
from pathlib import Path
from unittest.mock import patch


ROOT_DIR = Path(__file__).resolve().parents[1]
SPEC = importlib.util.spec_from_file_location("modelhub_agent_entrypoint", ROOT_DIR / "main.py")
if SPEC is None or SPEC.loader is None:
    raise RuntimeError("Unable to load the hosted agent entrypoint")
ENTRYPOINT = importlib.util.module_from_spec(SPEC)
SPEC.loader.exec_module(ENTRYPOINT)


class HostedAgentEntrypointTests(unittest.TestCase):
    def test_hosted_worker_always_uses_unlimited_cycle_submissions(self) -> None:
        with patch.dict(
            os.environ,
            {
                "MODELHUB_AGENT_MAX_SUBMITS_PER_RUN": "5",
                "MODELHUB_AGENT_EXTRA_ARGS": "--max-submits-per-run 3",
            },
            clear=False,
        ):
            command = ENTRYPOINT._worker_command()

        self.assertEqual(["--max-submits-per-run", "0"], command[-2:])
        self.assertIn("--state-sync", command)

    def test_readiness_file_is_separate_from_liveness(self) -> None:
        with tempfile.TemporaryDirectory() as temporary_dir:
            path = Path(temporary_dir) / "readiness.json"
            path.write_text('{"ready": false, "reason": "state_sync_unhealthy"}', encoding="utf-8")
            with patch.object(ENTRYPOINT, "READINESS_PATH", path):
                readiness = ENTRYPOINT._readiness()
        self.assertFalse(readiness["ready"])
        self.assertEqual("state_sync_unhealthy", readiness["reason"])

    def test_runtime_image_uses_python_git_client_without_apt_layer(self) -> None:
        dockerfile = (ROOT_DIR / "Dockerfile").read_text(encoding="utf-8")
        requirements = (ROOT_DIR / "requirements.txt").read_text(encoding="utf-8")
        self.assertNotIn("apt-get", dockerfile)
        self.assertIn("dulwich", requirements.casefold())

    def test_exited_worker_is_backed_off_without_exiting_supervisor(self) -> None:
        class ExitedWorker:
            def poll(self) -> int:
                return 137

        with tempfile.TemporaryDirectory() as temporary_dir:
            root = Path(temporary_dir)
            with (
                patch.object(ENTRYPOINT, "worker", ExitedWorker()),
                patch.object(ENTRYPOINT, "worker_start_enabled", True),
                patch.object(ENTRYPOINT, "shutdown_requested", False),
                patch.object(ENTRYPOINT, "config_error", None),
                patch.object(ENTRYPOINT, "worker_started_at", 90.0),
                patch.object(ENTRYPOINT, "worker_restart_count", 0),
                patch.object(ENTRYPOINT, "worker_next_restart_at", 0.0),
                patch.object(ENTRYPOINT, "WORKER_CRASH_PATH", root / "worker_crashes.jsonl"),
                patch.object(ENTRYPOINT, "READINESS_PATH", root / "readiness.json"),
            ):
                ENTRYPOINT._supervise_worker(now=100.0)
                self.assertIsNone(ENTRYPOINT.worker)
                self.assertEqual(1, ENTRYPOINT.worker_restart_count)
                self.assertEqual(105.0, ENTRYPOINT.worker_next_restart_at)
                self.assertEqual("worker_restarting", ENTRYPOINT._readiness()["reason"])
                self.assertIn('"exitCode":137', (root / "worker_crashes.jsonl").read_text(encoding="utf-8"))


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