ml-intern / tests /unit /test_cli_rendering.py
lewtun's picture
lewtun HF Staff
Add Slack gateway (#116)
6155b26 unverified
raw
history blame
2.81 kB
"""Regression tests for interactive CLI rendering and research model routing."""
import sys
from io import StringIO
from types import SimpleNamespace
import pytest
import agent.main as main_mod
from agent.tools.research_tool import _get_research_model
from agent.utils import terminal_display
def test_direct_anthropic_research_model_stays_off_bedrock():
assert _get_research_model("anthropic/claude-opus-4-6") == "anthropic/claude-sonnet-4-6"
def test_bedrock_anthropic_research_model_stays_on_bedrock():
assert (
_get_research_model("bedrock/us.anthropic.claude-opus-4-6-v1")
== "bedrock/us.anthropic.claude-sonnet-4-6"
)
def test_non_anthropic_research_model_is_unchanged():
assert _get_research_model("openai/gpt-5.4") == "openai/gpt-5.4"
def test_subagent_display_does_not_spawn_background_redraw(monkeypatch):
calls: list[object] = []
def _unexpected_future(*args, **kwargs):
calls.append((args, kwargs))
raise AssertionError("background redraw task should not be created")
monkeypatch.setattr("asyncio.ensure_future", _unexpected_future)
monkeypatch.setattr(
terminal_display,
"_console",
SimpleNamespace(file=StringIO(), width=100),
)
mgr = terminal_display.SubAgentDisplayManager()
mgr.start("agent-1", "research")
mgr.add_call("agent-1", "▸ hf_papers {\"operation\": \"search\"}")
mgr.clear("agent-1")
assert calls == []
def test_cli_forwards_model_flag_to_interactive_main(monkeypatch):
seen: dict[str, str | None] = {}
async def fake_main(*, model=None):
seen["model"] = model
monkeypatch.setattr(sys, "argv", ["ml-intern", "--model", "openai/gpt-5.5"])
monkeypatch.setattr(main_mod, "main", fake_main)
main_mod.cli()
assert seen["model"] == "openai/gpt-5.5"
@pytest.mark.asyncio
async def test_interactive_main_applies_model_override_before_banner(monkeypatch):
class StopAfterBanner(Exception):
pass
def fake_banner(*, model=None, hf_user=None):
assert model == "openai/gpt-5.5"
assert hf_user == "tester"
raise StopAfterBanner
monkeypatch.setattr(main_mod.os, "system", lambda *_args, **_kwargs: 0)
monkeypatch.setattr(main_mod, "PromptSession", lambda: object())
monkeypatch.setattr(main_mod, "resolve_hf_token", lambda: "hf-token")
monkeypatch.setattr(main_mod, "_get_hf_user", lambda _token: "tester")
monkeypatch.setattr(
main_mod,
"load_config",
lambda _path, **_kwargs: SimpleNamespace(
model_name="moonshotai/Kimi-K2.6",
mcpServers={},
),
)
monkeypatch.setattr(main_mod, "print_banner", fake_banner)
with pytest.raises(StopAfterBanner):
await main_mod.main(model="openai/gpt-5.5")