You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
86 lines
2.8 KiB
86 lines
2.8 KiB
"""
|
|
Unit tests for model configurations
|
|
"""
|
|
import pytest
|
|
from unittest.mock import patch, Mock
|
|
from src.models.openai import OpenAILikeModel
|
|
from src.models.openrouter import OpenRouterModel
|
|
|
|
|
|
class TestOpenAILikeModel:
|
|
"""Test cases for OpenAI-like model"""
|
|
|
|
@patch.dict('os.environ', {
|
|
'MODEL_ID': 'test-model',
|
|
'API_URL': 'https://test.api.com',
|
|
'MEGALLM_API_KEY': 'test-key'
|
|
})
|
|
def test_model_initialization_with_env(self):
|
|
"""Test model initialization with environment variables"""
|
|
model = OpenAILikeModel()
|
|
|
|
assert model.model_id == 'test-model'
|
|
assert model.api_url == 'https://test.api.com'
|
|
assert model.api_key == 'test-key'
|
|
|
|
def test_model_initialization_with_params(self):
|
|
"""Test model initialization with explicit parameters"""
|
|
model = OpenAILikeModel(
|
|
model_id='custom-model',
|
|
api_url='https://custom.api.com',
|
|
api_key='custom-key'
|
|
)
|
|
|
|
assert model.model_id == 'custom-model'
|
|
assert model.api_url == 'https://custom.api.com'
|
|
assert model.api_key == 'custom-key'
|
|
|
|
@patch('src.models.openai.OpenAILike')
|
|
def test_get_model(self, mock_openai_like):
|
|
"""Test get_model returns configured OpenAI-like model"""
|
|
mock_instance = Mock()
|
|
mock_openai_like.return_value = mock_instance
|
|
|
|
model = OpenAILikeModel()
|
|
result = model.get_model()
|
|
|
|
mock_openai_like.assert_called_once_with(
|
|
id=model.model_id,
|
|
api_key=model.api_key,
|
|
base_url=model.api_url,
|
|
default_headers={
|
|
"Authorization": f"Bearer {model.api_key}",
|
|
"Content-Type": "application/json"
|
|
}
|
|
)
|
|
assert result == mock_instance
|
|
|
|
|
|
class TestOpenRouterModel:
|
|
"""Test cases for OpenRouter model"""
|
|
|
|
def test_model_initialization_default(self):
|
|
"""Test model initialization with default values"""
|
|
model = OpenRouterModel()
|
|
|
|
assert model.model_id == "deepseek/deepseek-r1-0528:free"
|
|
assert model.api_key is None
|
|
|
|
def test_model_initialization_custom(self):
|
|
"""Test model initialization with custom values"""
|
|
model = OpenRouterModel(model_id="custom/model", api_key="custom-key")
|
|
|
|
assert model.model_id == "custom/model"
|
|
assert model.api_key == "custom-key"
|
|
|
|
@patch('src.models.openrouter.OpenRouter')
|
|
def test_get_model(self, mock_openrouter):
|
|
"""Test get_model returns configured OpenRouter model"""
|
|
mock_instance = Mock()
|
|
mock_openrouter.return_value = mock_instance
|
|
|
|
model = OpenRouterModel()
|
|
result = model.get_model()
|
|
|
|
mock_openrouter.assert_called_once_with(id=model.model_id)
|
|
assert result == mock_instance
|