test: add reasoning attribute handling to MockStreamingChunk

This commit is contained in:
Paul Gauthier (aider) 2025-03-09 08:40:26 -07:00
parent a15d10ea1e
commit 0df959cf68

View file

@ -90,7 +90,7 @@ class TestReasoning(unittest.TestCase):
# Mock streaming response chunks
class MockStreamingChunk:
def __init__(self, content=None, reasoning_content=None, finish_reason=None):
def __init__(self, content=None, reasoning_content=None, reasoning=None, finish_reason=None):
self.choices = [MagicMock()]
self.choices[0].delta = MagicMock()
self.choices[0].finish_reason = finish_reason
@ -109,6 +109,13 @@ class TestReasoning(unittest.TestCase):
# Need to handle attribute access that would raise AttributeError
delattr(self.choices[0].delta, "reasoning_content")
# Set reasoning if provided
if reasoning is not None:
self.choices[0].delta.reasoning = reasoning
else:
# Need to handle attribute access that would raise AttributeError
delattr(self.choices[0].delta, "reasoning")
# Create chunks to simulate streaming
chunks = [
# First chunk with reasoning content starts the tag
@ -264,7 +271,7 @@ class TestReasoning(unittest.TestCase):
# Mock streaming response chunks
class MockStreamingChunk:
def __init__(self, content=None, reasoning_content=None, finish_reason=None):
def __init__(self, content=None, reasoning_content=None, reasoning=None, finish_reason=None):
self.choices = [MagicMock()]
self.choices[0].delta = MagicMock()
self.choices[0].finish_reason = finish_reason
@ -283,6 +290,13 @@ class TestReasoning(unittest.TestCase):
# Need to handle attribute access that would raise AttributeError
delattr(self.choices[0].delta, "reasoning_content")
# Set reasoning if provided
if reasoning is not None:
self.choices[0].delta.reasoning = reasoning
else:
# Need to handle attribute access that would raise AttributeError
delattr(self.choices[0].delta, "reasoning")
# Create chunks to simulate streaming with think tags
chunks = [
# Start with open think tag