111 lines
4.0 KiB
Python
111 lines
4.0 KiB
Python
import os
|
||
from unittest.mock import MagicMock
|
||
|
||
import pytest
|
||
|
||
from app.services.ai_engine.base import AIEngineAdapter, EngineType
|
||
from app.services.api_key_manager import APIKeyManager, KeySource
|
||
|
||
|
||
class MockAdapter(AIEngineAdapter):
|
||
def __init__(self, api_key: str | None = None, **kwargs):
|
||
super().__init__(api_key=api_key, **kwargs)
|
||
|
||
async def query(self, query: str, brand_name: str, competitor_names: list[str] | None = None):
|
||
pass
|
||
|
||
def get_engine_type(self) -> EngineType:
|
||
return EngineType.CHATGPT
|
||
|
||
def _get_env_key(self) -> str | None:
|
||
return os.getenv("OPENAI_API_KEY", "")
|
||
|
||
|
||
class TestAdapterKeySource:
|
||
def test_adapter_accepts_key_manager_parameter(self):
|
||
key_manager = MagicMock(spec=APIKeyManager)
|
||
adapter = MockAdapter(key_manager=key_manager, user_id="user123")
|
||
assert adapter._key_manager is key_manager
|
||
assert adapter._user_id == "user123"
|
||
|
||
def test_adapter_accepts_api_key_parameter(self):
|
||
adapter = MockAdapter(api_key="direct-key-123")
|
||
assert adapter.api_key == "direct-key-123"
|
||
|
||
def test_resolve_key_from_direct_api_key(self):
|
||
adapter = MockAdapter(api_key="direct-key-123")
|
||
assert adapter.api_key == "direct-key-123"
|
||
|
||
def test_resolve_key_from_key_manager_user_key(self):
|
||
key_manager = MagicMock(spec=APIKeyManager)
|
||
key_manager.get_key.return_value = "manager-key-456"
|
||
|
||
adapter = MockAdapter(key_manager=key_manager, user_id="user123")
|
||
|
||
key_manager.get_key.assert_called_once_with("chatgpt", user_id="user123")
|
||
assert adapter.api_key == "manager-key-456"
|
||
|
||
def test_resolve_key_fallback_to_env_when_no_manager(self, monkeypatch):
|
||
monkeypatch.setenv("OPENAI_API_KEY", "env-key-789")
|
||
|
||
adapter = MockAdapter()
|
||
assert adapter.api_key == "env-key-789"
|
||
|
||
def test_direct_api_key_priority_over_key_manager(self):
|
||
key_manager = MagicMock(spec=APIKeyManager)
|
||
key_manager.get_key.return_value = "manager-key-456"
|
||
|
||
adapter = MockAdapter(api_key="direct-key-123", key_manager=key_manager, user_id="user123")
|
||
|
||
key_manager.get_key.assert_not_called()
|
||
assert adapter.api_key == "direct-key-123"
|
||
|
||
def test_user_key_priority_over_system_key(self):
|
||
key_manager = MagicMock(spec=APIKeyManager)
|
||
key_manager.get_key.return_value = "user-key-from-manager"
|
||
|
||
adapter = MockAdapter(key_manager=key_manager, user_id="user123")
|
||
assert adapter.api_key == "user-key-from-manager"
|
||
|
||
def test_no_key_available_returns_empty(self, monkeypatch):
|
||
key_manager = MagicMock(spec=APIKeyManager)
|
||
key_manager.get_key.return_value = None
|
||
|
||
# 清除环境变量，确保不会从环境变量获取key
|
||
monkeypatch.delenv("OPENAI_API_KEY", raising=False)
|
||
|
||
adapter = MockAdapter(key_manager=key_manager, user_id="user123")
|
||
assert adapter.api_key == ""
|
||
|
||
def test_adapter_with_all_parameters(self):
|
||
key_manager = MagicMock(spec=APIKeyManager)
|
||
key_manager.get_key.return_value = "full-key"
|
||
|
||
adapter = MockAdapter(
|
||
api_key=None,
|
||
rate_limiter=MagicMock(),
|
||
proxy="http://proxy:8080",
|
||
key_manager=key_manager,
|
||
user_id="test-user"
|
||
)
|
||
assert adapter._key_manager is key_manager
|
||
assert adapter._user_id == "test-user"
|
||
assert adapter.proxy == "http://proxy:8080"
|
||
|
||
|
||
class TestBatchQueryServiceKeyIntegration:
|
||
def test_build_adapters_with_key_manager(self):
|
||
from app.services.ai_engine.batch_query import BatchQueryService
|
||
|
||
key_manager = MagicMock(spec=APIKeyManager)
|
||
key_manager.get_key.return_value = "batch-test-key"
|
||
|
||
adapters = {
|
||
"chatgpt": MockAdapter(key_manager=key_manager, user_id="batch-user")
|
||
}
|
||
service = BatchQueryService(adapters)
|
||
service.set_user_context(user_id="batch-user", brand_id="brand-123")
|
||
|
||
assert service._user_id == "batch-user"
|
||
assert service._brand_id == "brand-123"
|