Initial Python SDK import

tests/test_client.py

@@ -0,0 +1,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)