97 lines
3.1 KiB
Python
97 lines
3.1 KiB
Python
import pytest
|
|
import httpx
|
|
import json
|
|
from typing import List, AsyncGenerator
|
|
|
|
from app.services import call_llm_api_real
|
|
from app.models import ChatMessage
|
|
from app.core.config import Settings
|
|
|
|
# Sample SSE chunks to simulate a streaming response
|
|
SSE_STREAM_CHUNKS = [
|
|
'data: {"choices": [{"delta": {"role": "assistant", "content": "Hello"}}]}',
|
|
'data: {"choices": [{"delta": {"content": " world!"}}]}',
|
|
'data: {"choices": [{"delta": {"tool_calls": [{"index": 0, "id": "call_123", "function": {"name": "get_weather", "arguments": ""}}]}}]}',
|
|
'data: {"choices": [{"delta": {"tool_calls": [{"index": 0, "function": {"arguments": "{\\"location\\":"}}]}}]}',
|
|
'data: {"choices": [{"delta": {"tool_calls": [{"index": 0, "function": {"arguments": " \\"San Francisco\\"}"}}]}}]}',
|
|
'data: [DONE]',
|
|
]
|
|
|
|
# Mock settings for the test
|
|
@pytest.fixture
|
|
def mock_settings() -> Settings:
|
|
"""Provides mock settings for tests."""
|
|
return Settings(
|
|
REAL_LLM_API_URL="http://fake-llm-api.com/chat",
|
|
REAL_LLM_API_KEY="fake-key"
|
|
)
|
|
|
|
# Async generator to mock the streaming response
|
|
async def mock_aiter_lines() -> AsyncGenerator[str, None]:
|
|
for chunk in SSE_STREAM_CHUNKS:
|
|
yield chunk
|
|
|
|
# Mock for the httpx.Response object
|
|
class MockStreamResponse:
|
|
def __init__(self, status_code: int = 200):
|
|
self._status_code = status_code
|
|
|
|
def raise_for_status(self):
|
|
if self._status_code != 200:
|
|
raise httpx.HTTPStatusError(
|
|
message="Error", request=httpx.Request("POST", ""), response=httpx.Response(self._status_code)
|
|
)
|
|
|
|
def aiter_lines(self) -> AsyncGenerator[str, None]:
|
|
return mock_aiter_lines()
|
|
|
|
async def __aenter__(self):
|
|
return self
|
|
|
|
async def __aexit__(self, exc_type, exc_val, exc_tb):
|
|
pass
|
|
|
|
# Mock for the httpx.AsyncClient
|
|
class MockAsyncClient:
|
|
def stream(self, method, url, headers, json, timeout):
|
|
return MockStreamResponse()
|
|
|
|
async def __aenter__(self):
|
|
return self
|
|
|
|
async def __aexit__(self, exc_type, exc_val, exc_tb):
|
|
pass
|
|
|
|
@pytest.mark.anyio
|
|
async def test_call_llm_api_real_streaming(monkeypatch, mock_settings):
|
|
"""
|
|
Tests that `call_llm_api_real` correctly handles an SSE stream,
|
|
parses the chunks, and assembles the final response message.
|
|
"""
|
|
# Patch httpx.AsyncClient to use our mock
|
|
monkeypatch.setattr(httpx, "AsyncClient", MockAsyncClient)
|
|
|
|
messages = [ChatMessage(role="user", content="What is the weather in San Francisco?")]
|
|
|
|
# Call the function
|
|
result = await call_llm_api_real(messages, mock_settings)
|
|
|
|
# Define the expected assembled result
|
|
expected_result = {
|
|
"role": "assistant",
|
|
"content": "Hello world!",
|
|
"tool_calls": [
|
|
{
|
|
"id": "call_123",
|
|
"type": "function",
|
|
"function": {
|
|
"name": "get_weather",
|
|
"arguments": '{"location": "San Francisco"}',
|
|
},
|
|
}
|
|
],
|
|
}
|
|
|
|
# Assert that the result matches the expected output
|
|
assert result == expected_result
|