|
| 1 | + |
| 2 | +import pytest |
| 3 | +from unittest.mock import MagicMock, Mock |
| 4 | +from supercoder.llm.openai_client import OpenAIClient |
| 5 | +from supercoder.llm.base import Message, StreamChunk |
| 6 | +from supercoder.config import Config |
| 7 | + |
| 8 | +class MockDelta: |
| 9 | + def __init__(self, content=None, reasoning_content=None): |
| 10 | + self.content = content |
| 11 | + if reasoning_content is not None: |
| 12 | + self.reasoning_content = reasoning_content |
| 13 | + |
| 14 | +class MockChoice: |
| 15 | + def __init__(self, delta): |
| 16 | + self.delta = delta |
| 17 | + |
| 18 | +class MockChunk: |
| 19 | + def __init__(self, choices=None): |
| 20 | + self.choices = choices |
| 21 | + |
| 22 | +def test_chat_stream_extracts_reasoning(): |
| 23 | + """Test that chat_stream extracts reasoning_content correctly.""" |
| 24 | + |
| 25 | + # Setup mock config |
| 26 | + config = Config() |
| 27 | + config.api_key = "test_key" |
| 28 | + |
| 29 | + # Setup mock client |
| 30 | + client = OpenAIClient(config) |
| 31 | + client.client = MagicMock() |
| 32 | + |
| 33 | + # Mock streaming response |
| 34 | + # 1. Chunk with reasoning only |
| 35 | + # 2. Chunk with content only |
| 36 | + # 3. Chunk with both |
| 37 | + mock_chunks = [ |
| 38 | + MockChunk(choices=[MockChoice(delta=MockDelta(content=None, reasoning_content="Thinking..."))]), |
| 39 | + MockChunk(choices=[MockChoice(delta=MockDelta(content="Hello", reasoning_content=None))]), |
| 40 | + MockChunk(choices=[MockChoice(delta=MockDelta(content=" World", reasoning_content=" (Done)"))]), |
| 41 | + ] |
| 42 | + |
| 43 | + client.client.chat.completions.create.return_value = mock_chunks |
| 44 | + |
| 45 | + # Run chat_stream |
| 46 | + chunks = list(client.chat_stream([Message("user", "test")])) |
| 47 | + |
| 48 | + # Verify chunks (ignoring the last "is_done" chunk) |
| 49 | + assert len(chunks) == 4 # 3 data chunks + 1 done chunk |
| 50 | + |
| 51 | + # Chunk 1: Reasoning only |
| 52 | + assert chunks[0].content == "" |
| 53 | + assert chunks[0].reasoning == "Thinking..." |
| 54 | + |
| 55 | + # Chunk 2: Content only |
| 56 | + assert chunks[1].content == "Hello" |
| 57 | + assert chunks[1].reasoning == "" |
| 58 | + |
| 59 | + # Chunk 3: Both |
| 60 | + assert chunks[2].content == " World" |
| 61 | + assert chunks[2].reasoning == " (Done)" |
| 62 | + |
| 63 | + # Chunk 4: Done |
| 64 | + assert chunks[3].is_done is True |
| 65 | + |
| 66 | +if __name__ == "__main__": |
| 67 | + test_chat_stream_extracts_reasoning() |
| 68 | + print("Test passed!") |
0 commit comments