diff --git a/python/tests/test_cli_flags.py b/python/tests/test_cli_flags.py index bae006f..6a01587 100644 --- a/python/tests/test_cli_flags.py +++ b/python/tests/test_cli_flags.py @@ -2,9 +2,11 @@ from __future__ import annotations +from unittest.mock import Mock + import pytest -from edge0.cli import _build_parser, _engine_kwargs +from edge0.cli import TIER_ENV, _build_parser, _engine_kwargs @pytest.mark.parametrize("cmd", ["demo", "chat", "serve"]) @@ -31,3 +33,31 @@ def test_serve_keeps_its_own_flags(): args = _build_parser().parse_args( ["serve", "edge0-8b", "--host", "0.0.0.0", "--port", "9"]) assert (args.host, args.port) == ("0.0.0.0", 9) + + +@pytest.mark.parametrize("cmd", ["demo", "chat", "serve"]) +@pytest.mark.parametrize("tier", ["edge0-8b", "edge0-35b"]) +@pytest.mark.parametrize("no_prerouter", [False, True]) +def test_entrypoint_preserves_prerouter_contract( + cmd, tier, no_prerouter, monkeypatch, tmp_path): + """Check actual CLI dispatch without constructing an inference engine.""" + from edge0 import AutoEngine + + class EngineConstructionReached(Exception): + pass + + factory = Mock(side_effect=EngineConstructionReached) + monkeypatch.setattr(AutoEngine, "from_pretrained", factory) + monkeypatch.setenv(TIER_ENV[tier], str(tmp_path)) + argv = [cmd, tier] + if no_prerouter: + argv.append("--no-prerouter") + args = _build_parser().parse_args(argv) + + with pytest.raises(EngineConstructionReached): + args.fn(args) + + expected = {} + if no_prerouter or (cmd == "demo" and tier == "edge0-8b"): + expected["prerouter"] = None + factory.assert_called_once_with(str(tmp_path), name=tier, **expected)