Initial public ModelForge release

node-agent/tests/test_main.py

@@ -0,0 +1,30 @@
+import sys
+
+import pytest
+
+import modelforge_node_agent.main as main_module
+from modelforge_node_agent.preflight import (
+    AgentStartupError,
+    AgentStartupFailureCode,
+    AgentStartupProblem,
+)
+
+
+def test_cli_reports_typed_preflight_failure_with_exit_three(monkeypatch) -> None:
+    async def fail_preflight(*, preflight_only: bool = False) -> None:
+        assert preflight_only
+        raise AgentStartupError(
+            AgentStartupProblem(
+                AgentStartupFailureCode.NVIDIA_NVML_UNAVAILABLE,
+                "MODELFORGE_AGENT_ACCELERATOR_MODE=nvidia",
+                "bounded diagnostic",
+            )
+        )
+
+    monkeypatch.setattr(main_module, "run_agent", fail_preflight)
+    monkeypatch.setattr(sys, "argv", ["modelforge-node-agent", "--preflight"])
+
+    with pytest.raises(SystemExit) as caught:
+        main_module.main()
+
+    assert caught.value.code == 3