From ebab75765dfb6fb827dfbb36d1191c5f19f7c3dc Mon Sep 17 00:00:00 2001 From: Paul Kompfner Date: Tue, 31 Mar 2026 18:54:23 -0400 Subject: [PATCH] Fix stream cancellation tests to mock get_chat_completions The tests were mocking the removed _stream_chat_completions_*_context methods. Update them to mock get_chat_completions instead. --- tests/test_novita_llm.py | 3 +-- tests/test_openai_llm_timeout.py | 6 ++---- tests/test_sambanova_llm.py | 3 +-- 3 files changed, 4 insertions(+), 8 deletions(-) diff --git a/tests/test_novita_llm.py b/tests/test_novita_llm.py index e0f2c71b4..9aab54cab 100644 --- a/tests/test_novita_llm.py +++ b/tests/test_novita_llm.py @@ -51,8 +51,7 @@ async def test_novita_llm_stream_closed_on_cancellation(): mock_stream = MockAsyncStream() - service._stream_chat_completions_specific_context = AsyncMock(return_value=mock_stream) - service._stream_chat_completions_universal_context = AsyncMock(return_value=mock_stream) + service.get_chat_completions = AsyncMock(return_value=mock_stream) service.start_ttfb_metrics = AsyncMock() service.stop_ttfb_metrics = AsyncMock() service.start_llm_usage_metrics = AsyncMock() diff --git a/tests/test_openai_llm_timeout.py b/tests/test_openai_llm_timeout.py index 8ee776a9b..61264cbf5 100644 --- a/tests/test_openai_llm_timeout.py +++ b/tests/test_openai_llm_timeout.py @@ -171,8 +171,7 @@ async def test_openai_llm_stream_closed_on_cancellation(): mock_stream = MockAsyncStream() # Mock the stream creation methods - service._stream_chat_completions_specific_context = AsyncMock(return_value=mock_stream) - service._stream_chat_completions_universal_context = AsyncMock(return_value=mock_stream) + service.get_chat_completions = AsyncMock(return_value=mock_stream) service.start_ttfb_metrics = AsyncMock() service.stop_ttfb_metrics = AsyncMock() service.start_llm_usage_metrics = AsyncMock() @@ -281,8 +280,7 @@ async def test_openai_llm_async_iterator_closed_on_stream_end(): mock_iterator = MockAsyncIterator() mock_stream = MockAsyncStream(mock_iterator) - service._stream_chat_completions_specific_context = AsyncMock(return_value=mock_stream) - service._stream_chat_completions_universal_context = AsyncMock(return_value=mock_stream) + service.get_chat_completions = AsyncMock(return_value=mock_stream) service.start_ttfb_metrics = AsyncMock() service.stop_ttfb_metrics = AsyncMock() service.start_llm_usage_metrics = AsyncMock() diff --git a/tests/test_sambanova_llm.py b/tests/test_sambanova_llm.py index 6632951fc..57e2e8c50 100644 --- a/tests/test_sambanova_llm.py +++ b/tests/test_sambanova_llm.py @@ -56,8 +56,7 @@ async def test_sambanova_llm_stream_closed_on_cancellation(): mock_stream = MockAsyncStream() - service._stream_chat_completions_specific_context = AsyncMock(return_value=mock_stream) - service._stream_chat_completions_universal_context = AsyncMock(return_value=mock_stream) + service.get_chat_completions = AsyncMock(return_value=mock_stream) service.start_ttfb_metrics = AsyncMock() service.stop_ttfb_metrics = AsyncMock() service.start_llm_usage_metrics = AsyncMock()