fix tests and add transport param to query

Signed-off-by: Rushil Patel <rpatel@codegen.com>

Author: rushilpatel0

src/claude_code_sdk/query.py

@@ -5,13 +5,15 @@
 from typing import Any
 
 from ._internal.client import InternalClient
+from ._internal.transport import Transport
 from .types import ClaudeCodeOptions, Message
 
 
 async def query(
     *,
     prompt: str | AsyncIterable[dict[str, Any]],
     options: ClaudeCodeOptions | None = None,
+    transport: Transport | None = None,
 ) -> AsyncIterator[Message]:
     """
     Query Claude Code for one-shot or unidirectional streaming interactions.
@@ -56,6 +58,9 @@ async def query(
                  - 'acceptEdits': Auto-accept file edits
                  - 'bypassPermissions': Allow all tools (use with caution)
                  Set options.cwd for working directory.
+        transport: Optional transport implementation. If provided, this will be used
+                  instead of the default transport selection based on options.
+                  The transport will be automatically configured with the prompt and options.
 
     Yields:
         Messages from the conversation
@@ -90,6 +95,22 @@ async def prompts():
         async for message in query(prompt=prompts()):
             print(message)
         ```
+
+    Example - With custom transport:
+        ```python
+        from claude_code_sdk import query, Transport
+        
+        class MyCustomTransport(Transport):
+            # Implement custom transport logic
+            pass
+        
+        transport = MyCustomTransport()
+        async for message in query(
+            prompt="Hello",
+            transport=transport
+        ):
+            print(message)
+        ```
     """
     if options is None:
         options = ClaudeCodeOptions()
@@ -98,5 +119,5 @@ async def prompts():
 
     client = InternalClient()
 
-    async for message in client.process_query(prompt=prompt, options=options):
+    async for message in client.process_query(prompt=prompt, options=options, transport=transport):
         yield message

tests/test_client.py

@@ -106,7 +106,9 @@ async def mock_receive():
                     messages.append(msg)
 
                 # Verify transport was created with correct parameters
-                mock_transport_class.assert_called_once_with()  # No parameters to constructor
-                mock_transport.configure.assert_called_once_with("test", options)
+                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)
+        anyio.run(_test)

tests/test_integration.py

@@ -195,10 +195,8 @@ async def mock_receive():
                     messages.append(msg)
 
                 # Verify transport was created with continuation option
-                mock_transport_class.assert_called_once_with()  # No parameters to constructor
-                mock_transport.configure.assert_called_once()
-                configure_call_args = mock_transport.configure.call_args
-                assert configure_call_args[0][0] == "Continue"  # prompt argument
-                assert configure_call_args[0][1].continue_conversation is True  # options argument
+                mock_transport_class.assert_called_once()
+                call_kwargs = mock_transport_class.call_args.kwargs
+                assert call_kwargs["options"].continue_conversation is True
 
-        anyio.run(_test)
+        anyio.run(_test)

tests/test_subprocess_buffering.py

@@ -50,8 +50,9 @@ async def _test() -> None:
 
             buffered_line = json.dumps(json_obj1) + "\n" + json.dumps(json_obj2)
 
-            transport = SubprocessCLITransport(cli_path="/usr/bin/claude")
-            transport.configure("test", ClaudeCodeOptions())
+            transport = SubprocessCLITransport(
+                prompt="test", options=ClaudeCodeOptions(), cli_path="/usr/bin/claude"
+            )
 
             mock_process = MagicMock()
             mock_process.returncode = None
@@ -84,8 +85,9 @@ async def _test() -> None:
 
             buffered_line = json.dumps(json_obj1) + "\n" + json.dumps(json_obj2)
 
-            transport = SubprocessCLITransport(cli_path="/usr/bin/claude")
-            transport.configure("test", ClaudeCodeOptions())
+            transport = SubprocessCLITransport(
+                prompt="test", options=ClaudeCodeOptions(), cli_path="/usr/bin/claude"
+            )
 
             mock_process = MagicMock()
             mock_process.returncode = None
@@ -113,8 +115,9 @@ async def _test() -> None:
 
             buffered_line = json.dumps(json_obj1) + "\n\n\n" + json.dumps(json_obj2)
 
-            transport = SubprocessCLITransport(cli_path="/usr/bin/claude")
-            transport.configure("test", ClaudeCodeOptions())
+            transport = SubprocessCLITransport(
+                prompt="test", options=ClaudeCodeOptions(), cli_path="/usr/bin/claude"
+            )
 
             mock_process = MagicMock()
             mock_process.returncode = None
@@ -158,8 +161,9 @@ async def _test() -> None:
             part2 = complete_json[100:250]
             part3 = complete_json[250:]
 
-            transport = SubprocessCLITransport(cli_path="/usr/bin/claude")
-            transport.configure("test", ClaudeCodeOptions())
+            transport = SubprocessCLITransport(
+                prompt="test", options=ClaudeCodeOptions(), cli_path="/usr/bin/claude"
+            )
 
             mock_process = MagicMock()
             mock_process.returncode = None
@@ -205,8 +209,9 @@ async def _test() -> None:
                 for i in range(0, len(complete_json), chunk_size)
             ]
 
-            transport = SubprocessCLITransport(cli_path="/usr/bin/claude")
-            transport.configure("test", ClaudeCodeOptions())
+            transport = SubprocessCLITransport(
+                prompt="test", options=ClaudeCodeOptions(), cli_path="/usr/bin/claude"
+            )
 
             mock_process = MagicMock()
             mock_process.returncode = None
@@ -234,8 +239,9 @@ def test_buffer_size_exceeded(self) -> None:
         async def _test() -> None:
             huge_incomplete = '{"data": "' + "x" * (_MAX_BUFFER_SIZE + 1000)
 
-            transport = SubprocessCLITransport(cli_path="/usr/bin/claude")
-            transport.configure("test", ClaudeCodeOptions())
+            transport = SubprocessCLITransport(
+                prompt="test", options=ClaudeCodeOptions(), cli_path="/usr/bin/claude"
+            )
 
             mock_process = MagicMock()
             mock_process.returncode = None
@@ -275,8 +281,9 @@ async def _test() -> None:
                 large_json[3000:] + "\n" + msg3,
             ]
 
-            transport = SubprocessCLITransport(cli_path="/usr/bin/claude")
-            transport.configure("test", ClaudeCodeOptions())
+            transport = SubprocessCLITransport(
+                prompt="test", options=ClaudeCodeOptions(), cli_path="/usr/bin/claude"
+            )
 
             mock_process = MagicMock()
             mock_process.returncode = None
@@ -297,4 +304,4 @@ async def _test() -> None:
             assert messages[2]["type"] == "system"
             assert messages[2]["subtype"] == "end"
 
-        anyio.run(_test)
+        anyio.run(_test)

tests/test_transport.py

@@ -21,14 +21,15 @@ def test_find_cli_not_found(self):
             patch("pathlib.Path.exists", return_value=False),
             pytest.raises(CLINotFoundError) as exc_info,
         ):
-            SubprocessCLITransport()
+            SubprocessCLITransport(prompt="test", options=ClaudeCodeOptions())
 
         assert "Claude Code requires Node.js" in str(exc_info.value)
 
     def test_build_command_basic(self):
         """Test building basic CLI command."""
-        transport = SubprocessCLITransport(cli_path="/usr/bin/claude")
-        transport.configure("Hello", ClaudeCodeOptions())
+        transport = SubprocessCLITransport(
+            prompt="Hello", options=ClaudeCodeOptions(), cli_path="/usr/bin/claude"
+        )
 
         cmd = transport._build_command()
         assert cmd[0] == "/usr/bin/claude"
@@ -41,23 +42,27 @@ def test_cli_path_accepts_pathlib_path(self):
         """Test that cli_path accepts pathlib.Path objects."""
         from pathlib import Path
 
-        transport = SubprocessCLITransport(cli_path=Path("/usr/bin/claude"))
+        transport = SubprocessCLITransport(
+            prompt="Hello",
+            options=ClaudeCodeOptions(),
+            cli_path=Path("/usr/bin/claude"),
+        )
 
         assert transport._cli_path == "/usr/bin/claude"
 
     def test_build_command_with_options(self):
         """Test building CLI command with options."""
-        transport = SubprocessCLITransport(cli_path="/usr/bin/claude")
-        transport.configure(
-            "test",
-            ClaudeCodeOptions(
+        transport = SubprocessCLITransport(
+            prompt="test",
+            options=ClaudeCodeOptions(
                 system_prompt="Be helpful",
                 allowed_tools=["Read", "Write"],
                 disallowed_tools=["Bash"],
                 model="claude-3-5-sonnet",
                 permission_mode="acceptEdits",
                 max_turns=5,
             ),
+            cli_path="/usr/bin/claude",
         )
 
         cmd = transport._build_command()
@@ -76,10 +81,10 @@ def test_build_command_with_options(self):
 
     def test_session_continuation(self):
         """Test session continuation options."""
-        transport = SubprocessCLITransport(cli_path="/usr/bin/claude")
-        transport.configure(
-            "Continue from before",
-            ClaudeCodeOptions(continue_conversation=True, resume="session-123"),
+        transport = SubprocessCLITransport(
+            prompt="Continue from before",
+            options=ClaudeCodeOptions(continue_conversation=True, resume="session-123"),
+            cli_path="/usr/bin/claude",
         )
 
         cmd = transport._build_command()
@@ -106,8 +111,11 @@ async def _test():
 
                 mock_exec.return_value = mock_process
 
-                transport = SubprocessCLITransport(cli_path="/usr/bin/claude")
-                transport.configure("test", ClaudeCodeOptions())
+                transport = SubprocessCLITransport(
+                    prompt="test",
+                    options=ClaudeCodeOptions(),
+                    cli_path="/usr/bin/claude",
+                )
 
                 await transport.connect()
                 assert transport._process is not None
@@ -122,8 +130,9 @@ def test_receive_messages(self):
         """Test parsing messages from CLI output."""
         # This test is simplified to just test the parsing logic
         # The full async stream handling is tested in integration tests
-        transport = SubprocessCLITransport(cli_path="/usr/bin/claude")
-        transport.configure("test", ClaudeCodeOptions())
+        transport = SubprocessCLITransport(
+            prompt="test", options=ClaudeCodeOptions(), cli_path="/usr/bin/claude"
+        )
 
         # The actual message parsing is done by the client, not the transport
         # So we just verify the transport can be created and basic structure is correct
@@ -135,26 +144,15 @@ def test_connect_with_nonexistent_cwd(self):
         from claude_code_sdk._errors import CLIConnectionError
 
         async def _test():
-            transport = SubprocessCLITransport(cli_path="/usr/bin/claude")
-            transport.configure(
-                "test",
-                ClaudeCodeOptions(cwd="/this/directory/does/not/exist"),
+            transport = SubprocessCLITransport(
+                prompt="test",
+                options=ClaudeCodeOptions(cwd="/this/directory/does/not/exist"),
+                cli_path="/usr/bin/claude",
             )
 
             with pytest.raises(CLIConnectionError) as exc_info:
                 await transport.connect()
 
             assert "/this/directory/does/not/exist" in str(exc_info.value)
 
-        anyio.run(_test)
-
-    def test_build_command_without_configure(self):
-        """Test that _build_command raises error if not configured."""
-        from claude_code_sdk._errors import CLIConnectionError
-
-        transport = SubprocessCLITransport(cli_path="/usr/bin/claude")
-        
-        with pytest.raises(CLIConnectionError) as exc_info:
-            transport._build_command()
-        
-        assert "Transport not configured" in str(exc_info.value)
+        anyio.run(_test)