mirrors/claude-code-sdk-python
1
0
Fork 0
mirror of https://github.com/anthropics/claude-code-sdk-python.git synced 2025-12-23 09:19:52 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Fix failing tests for control protocol implementation

Browse source 
- Update test mocks to handle control protocol initialization
- Fix interrupt test to check for control requests via write() instead of interrupt()
- Add control protocol support to mock_receive functions in streaming tests
- Ensure all tests properly handle init handshake before sending messages
- Fix formatting issues (trailing whitespace, blank lines)

All tests now properly support the new Query/Transport architecture with
bidirectional control protocol communication.
This commit is contained in:
Kashyap Murali 2025-09-01 03:24:02 -07:00
parent 105da74f68
commit a50b134ecf
No known key found for this signature in database
1 changed files with 71 additions and 2 deletions
Download patch file Download diff file Expand all files Collapse all files

73
tests/test_streaming_client.py
View file

@ -73,6 +73,33 @@ def create_mock_transport(with_init_response=True):
except (json.JSONDecodeError, KeyError, AttributeError):
pass
# Keep checking for other control requests (like interrupt)
last_check = len(written_messages)
timeout_counter = 0
while timeout_counter < 100: # Avoid infinite loop
await asyncio.sleep(0.01)
timeout_counter += 1
# Check for new messages
for msg_str in written_messages[last_check:]:
try:
msg = json.loads(msg_str.strip())
if msg.get("type") == "control_request":
subtype = msg.get("request", {}).get("subtype")
if subtype == "interrupt":
# Send interrupt response
yield {
"type": "control_response",
"response": {
"request_id": msg.get("request_id"),
"subtype": "success",
},
}
return # End after interrupt
except (json.JSONDecodeError, KeyError, AttributeError):
pass
last_check = len(written_messages)
# Then end the stream
return
@ -326,8 +353,34 @@ class TestClaudeSDKClientStreaming:
mock_transport = create_mock_transport()
mock_transport_class.return_value = mock_transport
# Mock the message stream
# Mock the message stream with control protocol support
async def mock_receive():
# First handle initialization
await asyncio.sleep(0.01)
written = mock_transport.write.call_args_list
for call in written:
data = call[0][0]
try:
msg = json.loads(data.strip())
if (
msg.get("type") == "control_request"
and msg.get("request", {}).get("subtype")
== "initialize"
):
yield {
"type": "control_response",
"response": {
"request_id": msg.get("request_id"),
"subtype": "success",
"commands": [],
"output_style": "default",
},
}
break
except (json.JSONDecodeError, KeyError, AttributeError):
pass
# Then yield the actual messages
yield {
"type": "assistant",
"message": {
@ -383,8 +436,24 @@ class TestClaudeSDKClientStreaming:
mock_transport_class.return_value = mock_transport
async with ClaudeSDKClient() as client:
# Interrupt is now handled via control protocol
await client.interrupt()
mock_transport.interrupt.assert_called_once()
# Check that a control request was sent via write
write_calls = mock_transport.write.call_args_list
interrupt_found = False
for call in write_calls:
data = call[0][0]
try:
msg = json.loads(data.strip())
if (
msg.get("type") == "control_request"
and msg.get("request", {}).get("subtype") == "interrupt"
):
interrupt_found = True
break
except (json.JSONDecodeError, KeyError, AttributeError):
pass
assert interrupt_found, "Interrupt control request not found"
anyio.run(_test)