claude-code-sdk-python/tests/test_client.py
yokomotod f1e6dda230
Some checks are pending
Lint / lint (push) Waiting to run
Test / test (3.10) (push) Waiting to run
Test / test (3.11) (push) Waiting to run
Test / test (3.12) (push) Waiting to run
Test / test (3.13) (push) Waiting to run
fix: add support for thinking content blocks (#28)
## Summary

Fixes an issue where `thinking` content blocks in Claude Code responses
were not being parsed, resulting in empty `AssistantMessage` content
arrays.

## Changes
- Added `ThinkingBlock` dataclass to handle thinking content with
`thinking` and `signature` fields
- Updated client parsing logic in `_internal/client.py` to recognize and
create `ThinkingBlock` instances
- Added comprehensive test coverage for thinking block functionality

## Before

```python
# Claude Code response with thinking block resulted in:
AssistantMessage(content=[])  # Empty content!
```

## After

```python
# Now correctly parses to:
AssistantMessage(content=[
    ThinkingBlock(thinking="...", signature="...")
])
```

Fixes #27

---------

Co-authored-by: Dickson Tsai <dickson@anthropic.com>
2025-08-06 19:00:02 -07:00

120 lines
4.3 KiB
Python

"""Tests for Claude SDK client functionality."""
from unittest.mock import AsyncMock, patch
import anyio
from claude_code_sdk import AssistantMessage, ClaudeCodeOptions, query
from claude_code_sdk.types import TextBlock
class TestQueryFunction:
"""Test the main query function."""
def test_query_single_prompt(self):
"""Test query with a single prompt."""
async def _test():
with patch(
"claude_code_sdk._internal.client.InternalClient.process_query"
) as mock_process:
# Mock the async generator
async def mock_generator():
yield AssistantMessage(
content=[TextBlock(text="4")], model="claude-opus-4-1-20250805"
)
mock_process.return_value = mock_generator()
messages = []
async for msg in query(prompt="What is 2+2?"):
messages.append(msg)
assert len(messages) == 1
assert isinstance(messages[0], AssistantMessage)
assert messages[0].content[0].text == "4"
anyio.run(_test)
def test_query_with_options(self):
"""Test query with various options."""
async def _test():
with patch(
"claude_code_sdk._internal.client.InternalClient.process_query"
) as mock_process:
async def mock_generator():
yield AssistantMessage(
content=[TextBlock(text="Hello!")],
model="claude-opus-4-1-20250805",
)
mock_process.return_value = mock_generator()
options = ClaudeCodeOptions(
allowed_tools=["Read", "Write"],
system_prompt="You are helpful",
permission_mode="acceptEdits",
max_turns=5,
)
messages = []
async for msg in query(prompt="Hi", options=options):
messages.append(msg)
# Verify process_query was called with correct prompt and options
mock_process.assert_called_once()
call_args = mock_process.call_args
assert call_args[1]["prompt"] == "Hi"
assert call_args[1]["options"] == options
anyio.run(_test)
def test_query_with_cwd(self):
"""Test query with custom working directory."""
async def _test():
with patch(
"claude_code_sdk._internal.client.SubprocessCLITransport"
) as mock_transport_class:
mock_transport = AsyncMock()
mock_transport_class.return_value = mock_transport
# Mock the message stream
async def mock_receive():
yield {
"type": "assistant",
"message": {
"role": "assistant",
"content": [{"type": "text", "text": "Done"}],
"model": "claude-opus-4-1-20250805",
},
}
yield {
"type": "result",
"subtype": "success",
"duration_ms": 1000,
"duration_api_ms": 800,
"is_error": False,
"num_turns": 1,
"session_id": "test-session",
"total_cost_usd": 0.001,
}
mock_transport.receive_messages = mock_receive
mock_transport.connect = AsyncMock()
mock_transport.disconnect = AsyncMock()
options = ClaudeCodeOptions(cwd="/custom/path")
messages = []
async for msg in query(prompt="test", options=options):
messages.append(msg)
# Verify transport was created with correct parameters
mock_transport_class.assert_called_once()
call_kwargs = mock_transport_class.call_args.kwargs
assert call_kwargs["prompt"] == "test"
assert call_kwargs["options"].cwd == "/custom/path"
anyio.run(_test)