Spaces:
Runtime error
Runtime error
File size: 1,313 Bytes
7ab7df1 ed005f8 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 | """Tests for models module."""
from unittest.mock import MagicMock, patch
import pytest
from models import create_embeddings, create_llm
@patch("models.ChatOpenAI")
def test_create_llm(mock_chat_openai):
"""Test LLM creation."""
mock_instance = MagicMock()
mock_chat_openai.return_value = mock_instance
result = create_llm(streaming=False)
mock_chat_openai.assert_called_once()
assert result == mock_instance
@patch("models.ChatOpenAI")
def test_create_llm_streaming(mock_chat_openai):
"""Test LLM creation with streaming enabled."""
mock_instance = MagicMock()
mock_chat_openai.return_value = mock_instance
result = create_llm(streaming=True)
call_kwargs = mock_chat_openai.call_args[1]
assert call_kwargs["streaming"] is True
assert result == mock_instance
@patch("models.HuggingFaceEmbeddings")
def test_create_embeddings(mock_embeddings):
"""Test embeddings creation."""
mock_instance = MagicMock()
mock_embeddings.return_value = mock_instance
result = create_embeddings()
mock_embeddings.assert_called_once()
call_kwargs = mock_embeddings.call_args[1]
assert call_kwargs["model_kwargs"]["device"] == "cpu"
assert call_kwargs["encode_kwargs"]["normalize_embeddings"] is True
assert result == mock_instance
|