012-kaopeilian/backend/tests/test_coze_client.py

"""
Coze 客户端单元测试
"""

import os
import pytest
from unittest.mock import Mock, patch, MagicMock
from cozepy import Coze, TokenAuth, OAuthJWT

from app.services.ai.coze.client import (
    CozeAuthManager, get_coze_client, get_bot_config,
    get_workspace_id
)
from app.services.ai.coze.exceptions import CozeAuthError


class TestCozeAuthManager:
    """测试认证管理器"""
    
    def test_init_with_env_vars(self):
        """测试从环境变量初始化"""
        with patch.dict(os.environ, {
            "COZE_API_BASE": "https://test.coze.cn",
            "COZE_WORKSPACE_ID": "test-workspace",
            "COZE_API_TOKEN": "test-token"
        }):
            manager = CozeAuthManager()
            assert manager.api_base == "https://test.coze.cn"
            assert manager.workspace_id == "test-workspace"
            assert manager.api_token == "test-token"
    
    def test_init_with_params(self):
        """测试从参数初始化"""
        manager = CozeAuthManager(
            api_base="https://custom.coze.cn",
            workspace_id="custom-workspace",
            api_token="custom-token"
        )
        assert manager.api_base == "https://custom.coze.cn"
        assert manager.workspace_id == "custom-workspace"
        assert manager.api_token == "custom-token"
    
    def test_setup_direct_connection(self):
        """测试直连设置"""
        manager = CozeAuthManager()
        no_proxy = os.environ.get("NO_PROXY", "")
        assert "api.coze.cn" in no_proxy
        assert ".coze.cn" in no_proxy
        assert "localhost" in no_proxy
    
    @patch("app.services.ai.coze.client.TokenAuth")
    @patch("app.services.ai.coze.client.Coze")
    def test_token_auth_success(self, mock_coze_class, mock_token_auth):
        """测试 Token 认证成功"""
        manager = CozeAuthManager(api_token="test-token")
        mock_client = Mock()
        mock_coze_class.return_value = mock_client
        
        client = manager.get_client()
        
        mock_token_auth.assert_called_once_with("test-token")
        mock_coze_class.assert_called_once()
        assert client == mock_client
    
    def test_token_auth_no_token(self):
        """测试没有 Token 时的错误"""
        manager = CozeAuthManager(api_token=None)
        
        with pytest.raises(CozeAuthError, match="API Token 未配置"):
            manager.get_client()
    
    @patch("builtins.open", create=True)
    @patch("app.services.ai.coze.client.serialization.load_pem_private_key")
    @patch("app.services.ai.coze.client.OAuthJWT")
    @patch("app.services.ai.coze.client.Coze")
    def test_oauth_auth_success(self, mock_coze_class, mock_oauth_jwt, 
                                mock_load_key, mock_open):
        """测试 OAuth 认证成功"""
        # 模拟私钥文件
        mock_open.return_value.__enter__.return_value.read.return_value = b"fake-private-key"
        mock_load_key.return_value = Mock()
        
        manager = CozeAuthManager(
            oauth_client_id="test-client",
            oauth_public_key_id="test-key-id",
            oauth_private_key_path="/path/to/key.pem"
        )
        
        mock_client = Mock()
        mock_coze_class.return_value = mock_client
        
        client = manager.get_client()
        
        mock_oauth_jwt.assert_called_once()
        assert client == mock_client
    
    @patch("builtins.open", side_effect=FileNotFoundError)
    @patch("app.services.ai.coze.client.TokenAuth")
    @patch("app.services.ai.coze.client.Coze")
    def test_oauth_fallback_to_token(self, mock_coze_class, mock_token_auth, mock_open):
        """测试 OAuth 失败后回退到 Token"""
        manager = CozeAuthManager(
            api_token="fallback-token",
            oauth_client_id="test-client",
            oauth_public_key_id="test-key-id",
            oauth_private_key_path="/nonexistent/key.pem"
        )
        
        mock_client = Mock()
        mock_coze_class.return_value = mock_client
        
        client = manager.get_client()
        
        # 应该使用 Token 认证
        mock_token_auth.assert_called_once_with("fallback-token")
        assert client == mock_client
    
    def test_refresh_token(self):
        """测试刷新令牌"""
        manager = CozeAuthManager(api_token="test-token")
        
        with patch.object(manager, '_init_client') as mock_init:
            manager.refresh_token()
            assert manager._client is None
            mock_init.assert_called_once()


class TestHelperFunctions:
    """测试辅助函数"""
    
    def test_get_bot_config(self):
        """测试获取 Bot 配置"""
        with patch.dict(os.environ, {
            "COZE_CHAT_BOT_ID": "chat-bot-123",
            "COZE_TRAINING_BOT_ID": "training-bot-456",
            "COZE_EXAM_BOT_ID": "exam-bot-789"
        }):
            config = get_bot_config()
            assert config["course_chat"] == "chat-bot-123"
            assert config["training"] == "training-bot-456"
            assert config["exam"] == "exam-bot-789"
    
    def test_get_workspace_id_success(self):
        """测试获取工作空间 ID 成功"""
        with patch.dict(os.environ, {"COZE_WORKSPACE_ID": "workspace-123"}):
            workspace_id = get_workspace_id()
            assert workspace_id == "workspace-123"
    
    def test_get_workspace_id_not_configured(self):
        """测试工作空间 ID 未配置"""
        with patch.dict(os.environ, {}, clear=True):
            with pytest.raises(CozeAuthError, match="COZE_WORKSPACE_ID 未配置"):
                get_workspace_id()
    
    @patch("app.services.ai.coze.client.get_auth_manager")
    def test_get_coze_client(self, mock_get_auth_manager):
        """测试获取 Coze 客户端"""
        mock_manager = Mock()
        mock_client = Mock()
        mock_manager.get_client.return_value = mock_client
        mock_get_auth_manager.return_value = mock_manager
        
        # 清除缓存
        get_coze_client.cache_clear()
        
        client = get_coze_client()
        assert client == mock_client
        mock_manager.get_client.assert_called_once()