56 lines
2.2 KiB
Python
56 lines
2.2 KiB
Python
import pytest
|
|
from unittest.mock import MagicMock, patch
|
|
from langchain_core.chat_history import InMemoryChatMessageHistory
|
|
from langchain_gigachat import GigaChat
|
|
|
|
from src.integrations.gigachat_api_client import GigaChatClient
|
|
from src.core.configuration import config
|
|
|
|
|
|
API_GIGACHAT_TOKEN: str = config.API_GIGACHAT_TOKEN
|
|
|
|
@pytest.fixture
|
|
def gigachat_client() -> GigaChatClient:
|
|
"""Fixture to create a GigaChatClient instance with a mock API token."""
|
|
return GigaChatClient(api_token=API_GIGACHAT_TOKEN)
|
|
|
|
def test_initialization(gigachat_client) -> None:
|
|
"""Test if the GigaChatClient initializes correctly."""
|
|
assert gigachat_client.api_token == API_GIGACHAT_TOKEN
|
|
assert gigachat_client.model_name == "GigaChat"
|
|
assert isinstance(gigachat_client.store, dict)
|
|
assert gigachat_client.llm is not None
|
|
|
|
def test_create_llm() -> None:
|
|
"""Test the _create_llm method for proper LLM creation."""
|
|
client = GigaChatClient(api_token=API_GIGACHAT_TOKEN)
|
|
mock_llm: GigaChat = client._create_llm("GigaChat-Pro")
|
|
assert mock_llm.credentials == API_GIGACHAT_TOKEN
|
|
assert mock_llm.model == "GigaChat-Pro"
|
|
|
|
def test_get_session_history(gigachat_client) -> None:
|
|
"""Test the get_session_history method for creating/retrieving session history."""
|
|
session_id = "test_session"
|
|
history = gigachat_client.get_session_history(session_id)
|
|
assert isinstance(history, InMemoryChatMessageHistory)
|
|
assert session_id in gigachat_client.store
|
|
assert gigachat_client.store[session_id] is history
|
|
|
|
def test_set_model(gigachat_client) -> None:
|
|
"""Test the set_model method for updating the LLM model."""
|
|
new_model = "GigaChat-Pro"
|
|
gigachat_client.set_model(new_model)
|
|
assert gigachat_client.llm.model == new_model
|
|
|
|
def test_get_response() -> None:
|
|
"""Test the get_response method by verifying the response code."""
|
|
with patch("langchain_core.runnables.history.RunnableWithMessageHistory") as MockRunnable:
|
|
mock_runnable = MagicMock()
|
|
mock_runnable.invoke.return_value.code = 200
|
|
MockRunnable.return_value = mock_runnable
|
|
|
|
client = GigaChatClient(api_token=API_GIGACHAT_TOKEN)
|
|
response_code = mock_runnable.invoke.return_value.code
|
|
|
|
assert response_code == 200
|