mirror of
https://github.com/anthropics/claude-code-sdk-python.git
synced 2025-07-07 14:45:00 +00:00
b3448f4baa
- Remove cost_usd field from test expectations - Update to use total_cost_usd instead of total_cost - All 30 tests now pass 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <noreply@anthropic.com>
114 lines
4 KiB
Python
114 lines
4 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")])
|
|
|
|
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!")])
|
|
|
|
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"}],
|
|
},
|
|
}
|
|
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)
|