forked from anthropics/claude-agent-sdk-python
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathtest_client.py
More file actions
115 lines (89 loc) · 4.07 KB
/
Copy pathtest_client.py
File metadata and controls
115 lines (89 loc) · 4.07 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
"""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",
"cost_usd": 0.001,
"duration_ms": 1000,
"duration_api_ms": 800,
"is_error": False,
"num_turns": 1,
"session_id": "test-session",
"total_cost": 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)