pyflowx/tests/cli/test_llm.py

"""Tests for ops.msdownload / ops.sglang 模块 (ModelScope 下载/SGLang 服务)."""

from __future__ import annotations

import subprocess
from pathlib import Path
from unittest.mock import MagicMock

import pytest

from pyflowx.conditions import Constants
from pyflowx.ops.msdownload import msdownload_run
from pyflowx.ops.sglang import install_sglang, run_sglang


# ---------------------------------------------------------------------- #
# msdownload_run
# ---------------------------------------------------------------------- #
class TestMsdownloadRun:
    """``msdownload_run`` 函数测试."""

    def test_empty_name_does_nothing(self, monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str]) -> None:
        """name 为空时应直接返回."""
        called: list[list[str]] = []
        monkeypatch.setattr(subprocess, "run", lambda cmd, **_: called.append(cmd))
        msdownload_run("")
        captured = capsys.readouterr()
        assert "name 不能为空" in captured.out
        assert called == []

    def test_default_download_dir(self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None:
        """未提供 download_dir 时默认使用 ~/.models/<name 最后一段>."""
        monkeypatch.setattr(Path, "home", lambda: tmp_path)
        ran_cmds: list[list[str]] = []
        monkeypatch.setattr(subprocess, "run", lambda cmd, **_: ran_cmds.append(cmd) or MagicMock())

        msdownload_run("Qwen/Qwen2.5-Coder")

        expected_dir = tmp_path / ".models" / "Qwen2.5-Coder"
        assert expected_dir.exists()
        assert ran_cmds[0] == [
            "uvx",
            "modelscope",
            "download",
            "--model",
            "Qwen/Qwen2.5-Coder",
            "--local_dir",
            str(expected_dir),
        ]

    def test_custom_download_dir(self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None:
        """提供 download_dir 时应使用指定目录."""
        custom_dir = tmp_path / "custom"
        ran_cmds: list[list[str]] = []
        monkeypatch.setattr(subprocess, "run", lambda cmd, **_: ran_cmds.append(cmd) or MagicMock())

        msdownload_run("Qwen/Qwen2.5", "dataset", str(custom_dir))

        assert custom_dir.exists()
        assert ran_cmds[0][3] == "--dataset"
        assert str(custom_dir) in ran_cmds[0]

    def test_dataset_type(self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None:
        """target_type=dataset 时应传递 --dataset."""
        monkeypatch.setattr(Path, "home", lambda: tmp_path)
        ran_cmds: list[list[str]] = []
        monkeypatch.setattr(subprocess, "run", lambda cmd, **_: ran_cmds.append(cmd) or MagicMock())

        msdownload_run("foo/bar", "dataset")

        assert "--dataset" in ran_cmds[0]


# ---------------------------------------------------------------------- #
# install_sglang
# ---------------------------------------------------------------------- #
class TestInstallSglang:
    """``install_sglang`` 函数测试."""

    def test_already_installed_skips(self, monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str]) -> None:
        """sglang 已安装时应跳过."""
        monkeypatch.setattr("shutil.which", lambda _cmd: "/usr/bin/sglang")
        called: list[list[str]] = []
        monkeypatch.setattr(subprocess, "run", lambda cmd, **_: called.append(cmd))

        install_sglang()

        captured = capsys.readouterr()
        assert "已安装" in captured.out
        assert called == []

    def test_not_installed_runs_uv_install(self, monkeypatch: pytest.MonkeyPatch) -> None:
        """sglang 未安装时应执行 uv install sglang[all]."""
        monkeypatch.setattr("shutil.which", lambda _cmd: None)
        ran_cmds: list[list[str]] = []
        monkeypatch.setattr(subprocess, "run", lambda cmd, **_: ran_cmds.append(cmd) or MagicMock())

        install_sglang()

        assert ran_cmds == [["uv", "install", "sglang[all]"]]


# ---------------------------------------------------------------------- #
# run_sglang
# ---------------------------------------------------------------------- #
class TestRunSglang:
    """``run_sglang`` 函数测试."""

    def test_model_dir_not_exist_skips(
        self, monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str]
    ) -> None:
        """模型目录不存在时应跳过."""
        monkeypatch.setattr(Path, "exists", lambda _self: False)
        called: list[list[str]] = []
        monkeypatch.setattr(subprocess, "run", lambda cmd, **_: called.append(cmd))

        run_sglang(model="/nonexistent/path")

        captured = capsys.readouterr()
        assert "模型目录不存在" in captured.out
        assert called == []

    def test_windows_uses_python(self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None:
        """Windows 上应使用 python."""
        monkeypatch.setattr(Constants, "IS_WINDOWS", True)
        monkeypatch.setattr(Path, "expanduser", lambda _self: tmp_path)
        monkeypatch.setattr(Path, "exists", lambda _self: True)
        ran_cmds: list[list[str]] = []
        monkeypatch.setattr(subprocess, "run", lambda cmd, **_: ran_cmds.append(cmd) or MagicMock())

        run_sglang(model=str(tmp_path))

        assert ran_cmds[0][0] == "python"
        assert "-m" in ran_cmds[0]
        assert "sglang.launch_server" in ran_cmds[0]

    def test_non_windows_uses_python3(self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None:
        """非 Windows 上应使用 python3."""
        monkeypatch.setattr(Constants, "IS_WINDOWS", False)
        monkeypatch.setattr(Path, "expanduser", lambda _self: tmp_path)
        monkeypatch.setattr(Path, "exists", lambda _self: True)
        ran_cmds: list[list[str]] = []
        monkeypatch.setattr(subprocess, "run", lambda cmd, **_: ran_cmds.append(cmd) or MagicMock())

        run_sglang(model=str(tmp_path), port=9000, ctx_len=4096, mem_fraction=0.5, host="127.0.0.1", log_level="debug")

        cmd = ran_cmds[0]
        assert cmd[0] == "python3"
        assert "--port" in cmd
        assert "9000" in cmd
        assert "4096" in cmd
        assert "0.5" in cmd
        assert "127.0.0.1" in cmd
        assert "debug" in cmd

    def test_command_includes_qwen_parser(self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None:
        """命令应包含 --tool-call-parser qwen."""
        monkeypatch.setattr(Constants, "IS_WINDOWS", True)
        monkeypatch.setattr(Path, "expanduser", lambda _self: tmp_path)
        monkeypatch.setattr(Path, "exists", lambda _self: True)
        ran_cmds: list[list[str]] = []
        monkeypatch.setattr(subprocess, "run", lambda cmd, **_: ran_cmds.append(cmd) or MagicMock())

        run_sglang(model=str(tmp_path))

        cmd = ran_cmds[0]
        assert "--tool-call-parser" in cmd
        qwen_idx = cmd.index("--tool-call-parser")
        assert cmd[qwen_idx + 1] == "qwen"