Fix AssistantMessage mocking

Author: dicksontsai

tests/test_client.py

@@ -20,7 +20,7 @@ async def _test():
             ) as mock_process:
                 # Mock the async generator
                 async def mock_generator():
-                    yield AssistantMessage(content=[TextBlock(text="4")])
+                    yield AssistantMessage(content=[TextBlock(text="4")], model="claude-opus-4-1-20250805")
 
                 mock_process.return_value = mock_generator()
 
@@ -43,7 +43,7 @@ async def _test():
             ) as mock_process:
 
                 async def mock_generator():
-                    yield AssistantMessage(content=[TextBlock(text="Hello!")])
+                    yield AssistantMessage(content=[TextBlock(text="Hello!")], model="claude-opus-4-1-20250805")
 
                 mock_process.return_value = mock_generator()
 
@@ -83,6 +83,7 @@ async def mock_receive():
                         "message": {
                             "role": "assistant",
                             "content": [{"type": "text", "text": "Done"}],
+                            "model": "claude-opus-4-1-20250805",
                         },
                     }
                     yield {

tests/test_integration.py

@@ -38,6 +38,7 @@ async def mock_receive():
                         "message": {
                             "role": "assistant",
                             "content": [{"type": "text", "text": "2 + 2 equals 4"}],
+                            "model": "claude-opus-4-1-20250805",
                         },
                     }
                     yield {
@@ -103,6 +104,7 @@ async def mock_receive():
                                     "input": {"file_path": "/test.txt"},
                                 },
                             ],
+                            "model": "claude-opus-4-1-20250805",
                         },
                     }
                     yield {
@@ -179,6 +181,7 @@ async def mock_receive():
                                     "text": "Continuing from previous conversation",
                                 }
                             ],
+                            "model": "claude-opus-4-1-20250805",
                         },
                     }
 

tests/test_streaming_client.py

@@ -187,6 +187,7 @@ async def mock_receive():
                         "message": {
                             "role": "assistant",
                             "content": [{"type": "text", "text": "Hello!"}],
+                            "model": "claude-opus-4-1-20250805",
                         },
                     }
                     yield {
@@ -229,6 +230,7 @@ async def mock_receive():
                         "message": {
                             "role": "assistant",
                             "content": [{"type": "text", "text": "Answer"}],
+                            "model": "claude-opus-4-1-20250805",
                         },
                     }
                     yield {
@@ -250,6 +252,7 @@ async def mock_receive():
                                 {"type": "text", "text": "Should not see this"}
                             ],
                         },
+                        "model": "claude-opus-4-1-20250805",
                     }
 
                 mock_transport.receive_messages = mock_receive
@@ -335,6 +338,7 @@ async def mock_receive():
                         "message": {
                             "role": "assistant",
                             "content": [{"type": "text", "text": "Response 1"}],
+                            "model": "claude-opus-4-1-20250805",
                         },
                     }
                     await asyncio.sleep(0.1)
@@ -531,13 +535,15 @@ async def mock_receive():
                         "message": {
                             "role": "assistant",
                             "content": [{"type": "text", "text": "Hello"}],
+                            "model": "claude-opus-4-1-20250805",
                         },
                     }
                     yield {
                         "type": "assistant",
                         "message": {
                             "role": "assistant",
                             "content": [{"type": "text", "text": "World"}],
+                            "model": "claude-opus-4-1-20250805",
                         },
                     }
                     yield {

tests/test_types.py

@@ -25,14 +25,14 @@ def test_user_message_creation(self):
     def test_assistant_message_with_text(self):
         """Test creating an AssistantMessage with text content."""
         text_block = TextBlock(text="Hello, human!")
-        msg = AssistantMessage(content=[text_block])
+        msg = AssistantMessage(content=[text_block], model="claude-opus-4-1-20250805")
         assert len(msg.content) == 1
         assert msg.content[0].text == "Hello, human!"
 
     def test_assistant_message_with_thinking(self):
         """Test creating an AssistantMessage with thinking content."""
         thinking_block = ThinkingBlock(thinking="I'm thinking...", signature="sig-123")
-        msg = AssistantMessage(content=[thinking_block])
+        msg = AssistantMessage(content=[thinking_block], model="claude-opus-4-1-20250805")
         assert len(msg.content) == 1
         assert msg.content[0].thinking == "I'm thinking..."
         assert msg.content[0].signature == "sig-123"