35 lines
1.0 KiB
Python
35 lines
1.0 KiB
Python
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import importlib.util
|
||
|
|
import os
|
||
|
|
import unittest
|
||
|
|
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:])
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
unittest.main()
|