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()