File size: 1,828 Bytes
f6fdf6a afd6869 f6fdf6a |
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 49 50 51 52 |
import pytest
from unittest.mock import patch, AsyncMock
import httpx
from app.providers.transformers_provider import list_models, chat
@pytest.mark.asyncio
async def test_list_models_success():
"""Test successful model listing."""
mock_response = {"data": [{"id": "test-model"}]}
with patch('httpx.AsyncClient') as mock_client:
mock_response_obj = AsyncMock()
mock_response_obj.json.return_value = mock_response
mock_response_obj.raise_for_status.return_value = None
mock_client.return_value.__aenter__.return_value.get.return_value = mock_response_obj
result = await list_models()
assert result == mock_response
@pytest.mark.asyncio
async def test_chat_success():
"""Test successful chat completion."""
payload = {"model": "test", "messages": [{"role": "user", "content": "hello"}]}
mock_response = {"choices": [{"message": {"content": "hi"}}]}
with patch('httpx.AsyncClient') as mock_client:
mock_response_obj = AsyncMock()
mock_response_obj.json.return_value = mock_response
mock_response_obj.raise_for_status.return_value = None
mock_client.return_value.__aenter__.return_value.post.return_value = mock_response_obj
result = await chat(payload, stream=False)
assert result == mock_response
@pytest.mark.asyncio
async def test_chat_stream():
"""Test chat completion with streaming."""
payload = {"model": "test", "messages": [{"role": "user", "content": "hello"}]}
mock_stream = AsyncMock()
with patch('httpx.AsyncClient') as mock_client:
mock_client.return_value.__aenter__.return_value.stream.return_value = mock_stream
result = await chat(payload, stream=True)
assert result == mock_stream
|