108 lines
3.9 KiB
Python
108 lines
3.9 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):
|
|
key_manager = MagicMock(spec=APIKeyManager)
|
|
key_manager.get_key.return_value = None
|
|
|
|
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"
|