FastDeploy/tests/entrypoints/cli/test_openai.py

import argparse
import signal
import unittest
from unittest.mock import MagicMock, call, patch

from fastdeploy.entrypoints.cli.openai import (
    ChatCommand,
    CompleteCommand,
    _add_query_options,
    _interactive_cli,
    _register_signal_handlers,
    chat,
    cmd_init,
)


class TestOpenAICli(unittest.TestCase):

    @patch("fastdeploy.entrypoints.cli.openai.signal.signal")
    def test_register_signal_handlers(self, mock_signal):
        """测试信号处理器注册"""
        _register_signal_handlers()

        # 验证信号处理器被正确注册
        mock_signal.assert_has_calls([call(signal.SIGINT, unittest.mock.ANY), call(signal.SIGTSTP, unittest.mock.ANY)])

    @patch("fastdeploy.entrypoints.cli.openai.OpenAI")
    @patch("fastdeploy.entrypoints.cli.openai.os.environ.get")
    @patch("fastdeploy.entrypoints.cli.openai._register_signal_handlers")
    def test_interactive_cli_with_model_name(self, mock_register, mock_environ, mock_openai):
        """测试交互式CLI初始化（指定模型名）"""
        # 设置mock
        mock_environ.return_value = "test_api_key"
        mock_client = MagicMock()
        mock_openai.return_value = mock_client

        # 测试参数
        args = argparse.Namespace()
        args.url = "http://localhost:9904/v1"
        args.api_key = None
        args.model_name = "test-model"

        # 执行测试
        model_name, client = _interactive_cli(args)

        # 验证结果
        self.assertEqual(model_name, "test-model")
        self.assertEqual(client, mock_client)
        mock_openai.assert_called_once_with(api_key="test_api_key", base_url="http://localhost:9904/v1")
        mock_register.assert_called_once()

    @patch("fastdeploy.entrypoints.cli.openai.OpenAI")
    @patch("fastdeploy.entrypoints.cli.openai.os.environ.get")
    @patch("fastdeploy.entrypoints.cli.openai._register_signal_handlers")
    def test_interactive_cli_without_model_name(self, mock_register, mock_environ, mock_openai):
        """测试交互式CLI初始化（未指定模型名）"""
        # 设置mock
        mock_environ.return_value = "test_api_key"
        mock_client = MagicMock()
        mock_models = MagicMock()
        mock_models.data = [MagicMock(id="first-model")]
        mock_client.models.list.return_value = mock_models
        mock_openai.return_value = mock_client

        # 测试参数
        args = argparse.Namespace()
        args.url = "http://localhost:9904/v1"
        args.api_key = None
        args.model_name = None

        # 执行测试
        model_name, client = _interactive_cli(args)

        # 验证结果
        self.assertEqual(model_name, "first-model")
        self.assertEqual(client, mock_client)
        mock_client.models.list.assert_called_once()

    @patch("fastdeploy.entrypoints.cli.openai.input")
    @patch("fastdeploy.entrypoints.cli.openai.OpenAI")
    def test_chat_function(self, mock_openai, mock_input):
        """测试chat函数"""
        # 设置mock
        mock_client = MagicMock()
        mock_completion = MagicMock()
        mock_completion.choices = [MagicMock()]
        mock_completion.choices[0].message = MagicMock(content="Test response")
        mock_client.chat.completions.create.return_value = mock_completion

        # 模拟用户输入和EOF
        mock_input.side_effect = ["Hello", EOFError]

        # 执行测试
        chat("System prompt", "test-model", mock_client)

        # 验证API调用
        mock_client.chat.completions.create.assert_called_once()

    def test_add_query_options(self):
        """测试查询选项添加"""
        mock_parser = MagicMock()

        result = _add_query_options(mock_parser)

        # 验证parser方法被调用
        self.assertEqual(result, mock_parser)
        self.assertEqual(mock_parser.add_argument.call_count, 3)

    def test_cmd_init(self):
        """测试命令初始化"""
        commands = cmd_init()

        # 验证返回的命令列表
        self.assertEqual(len(commands), 2)
        self.assertEqual(commands[0].name, "chat")
        self.assertEqual(commands[1].name, "complete")


class TestChatCommand(unittest.TestCase):

    @patch("fastdeploy.entrypoints.cli.openai._interactive_cli")
    @patch("fastdeploy.entrypoints.cli.openai.OpenAI")
    def test_chat_command_quick_mode(self, mock_openai, mock_interactive):
        """测试ChatCommand快速模式"""
        # 设置mock
        mock_interactive.return_value = ("test-model", MagicMock())

        args = argparse.Namespace()
        args.quick = "Quick message"
        args.system_prompt = None

        # 执行测试
        ChatCommand.cmd(args)

        # 验证_interactive_cli被调用
        mock_interactive.assert_called_once_with(args)

    @patch("fastdeploy.entrypoints.cli.openai.input")
    @patch("fastdeploy.entrypoints.cli.openai.OpenAI")
    def test_chat_empty_input(self, mock_openai, mock_input):
        """Test chat with empty input."""
        mock_client = MagicMock()
        mock_openai.return_value = mock_client

        # Mock empty input then EOF
        mock_input.side_effect = ["", EOFError()]

        args = argparse.Namespace()
        args.quick = None
        args.url = "http://test.com"
        args.api_key = None

        args.model_name = None
        args.system_prompt = "System prompt"

        ChatCommand().cmd(args)


class TestCompleteCommand(unittest.TestCase):

    @patch("fastdeploy.entrypoints.cli.openai._interactive_cli")
    @patch("fastdeploy.entrypoints.cli.openai.OpenAI")
    def test_complete_command_quick_mode(self, mock_openai, mock_interactive):
        """测试CompleteCommand快速模式"""
        # 设置mock
        mock_client = MagicMock()
        mock_completion = MagicMock()
        mock_completion.choices = [MagicMock(text="Completion text")]
        mock_client.completions.create.return_value = mock_completion
        mock_interactive.return_value = ("test-model", mock_client)

        args = argparse.Namespace()
        args.quick = "Quick prompt"

        # 执行测试
        CompleteCommand.cmd(args)

        # 验证API调用
        mock_client.completions.create.assert_called_once_with(model="test-model", prompt="Quick prompt")

    @patch("fastdeploy.entrypoints.cli.openai.input")
    @patch("fastdeploy.entrypoints.cli.openai.OpenAI")
    def test_completion_empty_input(self, mock_openai, mock_input):
        """Test completion with empty input."""
        mock_client = MagicMock()
        mock_openai.return_value = mock_client

        # Mock empty input then EOF
        mock_input.side_effect = ["", EOFError()]

        args = argparse.Namespace()
        args.quick = None
        args.url = "http://test.com"
        args.api_key = None
        args.model_name = None

        CompleteCommand.cmd(args)


if __name__ == "__main__":
    unittest.main()