diff --git a/tests/conftest.py b/tests/conftest.py index d1d0f66ef..b32a7e1f2 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -6,7 +6,7 @@ import tempfile import time from pathlib import Path -from unittest.mock import AsyncMock +from unittest.mock import AsyncMock, MagicMock from uuid import uuid4 import pytest @@ -101,43 +101,19 @@ def mock_call_llm_api(mocker): completion_mock = mocker.patch.object(LlmApiHandler, "call_llm_api") def wrap_unstreamed_string(value): - timestamp = int(time.time()) - return ChatCompletion( - id="test-id", - choices=[ - Choice( - finish_reason="stop", - index=0, - message=ChatCompletionMessage( - content=value, - role="assistant", - ), - ) - ], - created=timestamp, - model="test-model", - object="chat.completion", - ) + mock_spice_response = MagicMock() + mock_spice_response.text = value def wrap_streamed_strings(values): async def _async_generator(): - timestamp = int(time.time()) for value in values: - yield ChatCompletionChunk( - id="test-id", - choices=[ - AsyncChoice( - delta=ChoiceDelta(content=value, role="assistant"), - finish_reason=None, - index=0, - ) - ], - created=timestamp, - model="test-model", - object="chat.completion.chunk", - ) - - return _async_generator() + yield value + + mock_spice_response = MagicMock() + mock_spice_response.stream = _async_generator + mock_spice_response.text = "".join(values) + + return mock_spice_response def set_streamed_values(values): completion_mock.return_value = wrap_streamed_strings(values)