tiny fix claude agent sdk test (#216)

Author: dinmukhamedm

tests/test_instrumentations/test_claude_agent/test_claude_sdk_client.py

@@ -1,23 +1,21 @@
-import asyncio
+import pytest
 
 from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
 from claude_agent_sdk import ClaudeSDKClient
 
 from mock_transport import MockClaudeTransport
 
 
-def test_claude_agent_query(span_exporter: InMemorySpanExporter):
-    async def _collect_messages():
-        async with ClaudeSDKClient(transport=MockClaudeTransport()) as client:
-            await client.query("What's the capital of France?")
-            async for _ in client.receive_response():
-                pass
+@pytest.mark.asyncio
+async def test_claude_agent_query(span_exporter: InMemorySpanExporter):
+    async with ClaudeSDKClient(transport=MockClaudeTransport()) as client:
+        await client.query("What's the capital of France?")
+        async for _ in client.receive_response():
+            pass
 
-            await client.query("What's the population of that city?")
-            async for _ in client.receive_response():
-                pass
-
-    asyncio.run(_collect_messages())
+        await client.query("What's the population of that city?")
+        async for _ in client.receive_response():
+            pass
 
     spans_tuple = span_exporter.get_finished_spans()
     spans = sorted(list(spans_tuple), key=lambda x: x.start_time)

tests/test_instrumentations/test_claude_agent/test_query.py

@@ -1,35 +1,27 @@
-import asyncio
+import claude_agent_sdk
+import pytest
 
 from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
 from claude_agent_sdk import ClaudeAgentOptions
-import claude_agent_sdk
 
 from mock_transport import MockClaudeTransport
 
-# Note: can not use alias "query" aka "from claude_agent_sdk import query" because it is not wrapped by Laminar
 
-
-def test_claude_agent_query(span_exporter: InMemorySpanExporter):
+@pytest.mark.asyncio
+async def test_claude_agent_query(span_exporter: InMemorySpanExporter):
     options = ClaudeAgentOptions(
         model="claude-sonnet-4-5",
         system_prompt="You are an expert software engineer.",
         permission_mode="acceptEdits",
     )
 
-    async def _collect_messages():
-        messages = []
-        async for message in claude_agent_sdk.query(
-            # prompt="Create a readme doc for the test_claude_agent package. Then delete it. Return task status.",
-            prompt="What is the capital of France?",
-            options=options,
-            transport=MockClaudeTransport(
-                auto_respond_on_connect=True, close_after_responses=True
-            ),
-        ):
-            messages.append(message)
-        return messages
-
-    for _ in asyncio.run(_collect_messages()):
+    async for message in claude_agent_sdk.query(
+        prompt="What is the capital of France?",
+        options=options,
+        transport=MockClaudeTransport(
+            auto_respond_on_connect=True, close_after_responses=True
+        ),
+    ):
         pass
 
     spans = span_exporter.get_finished_spans()

tests/test_instrumentations/test_claude_agent/test_query_with_alias.py

@@ -1,32 +1,26 @@
-import asyncio
+import pytest
 
 from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
 from claude_agent_sdk import ClaudeAgentOptions, query as claude_query
 
 from mock_transport import MockClaudeTransport
 
 
-def test_claude_agent_query(span_exporter: InMemorySpanExporter):
+@pytest.mark.asyncio
+async def test_claude_agent_query(span_exporter: InMemorySpanExporter):
     options = ClaudeAgentOptions(
         model="claude-sonnet-4-5",
         system_prompt="You are an expert software engineer.",
         permission_mode="acceptEdits",
     )
 
-    async def _collect_messages():
-        messages = []
-        async for message in claude_query(
-            # prompt="Create a readme doc for the test_claude_agent package. Then delete it. Return task status.",
-            prompt="What is the capital of France?",
-            options=options,
-            transport=MockClaudeTransport(
-                auto_respond_on_connect=True, close_after_responses=True
-            ),
-        ):
-            messages.append(message)
-        return messages
-
-    for _ in asyncio.run(_collect_messages()):
+    async for message in claude_query(
+        prompt="What is the capital of France?",
+        options=options,
+        transport=MockClaudeTransport(
+            auto_respond_on_connect=True, close_after_responses=True
+        ),
+    ):
         pass
 
     spans = span_exporter.get_finished_spans()

tests/test_instrumentations/test_claude_agent/test_tool.py

@@ -1,4 +1,4 @@
-import asyncio
+import pytest
 from typing import Any
 from claude_agent_sdk import (
     tool,
@@ -31,31 +31,27 @@ async def get_time(args: dict[str, Any]) -> dict[str, Any]:
     return {"content": [{"type": "text", "text": f"Current time: {current_time}"}]}
 
 
-def test_claude_agent_tool(span_exporter: InMemorySpanExporter):
+@pytest.mark.asyncio
+async def test_claude_agent_tool(span_exporter: InMemorySpanExporter):
     my_server = create_sdk_mcp_server(
         name="utilities", version="1.0.0", tools=[calculate, get_time]
     )
 
-    # Configure options with the server
     options = ClaudeAgentOptions(
         mcp_servers={"utils": my_server},
         allowed_tools=["mcp__utils__calculate", "mcp__utils__get_time"],
     )
 
-    # Use ClaudeSDKClient for interactive tool usage
-    async def _collect_messages():
-        async with ClaudeSDKClient(
-            options=options, transport=MockClaudeTransport()
-        ) as client:
-            await client.query("What's 123 * 456?")
-            async for _ in client.receive_response():
-                pass
+    async with ClaudeSDKClient(
+        options=options, transport=MockClaudeTransport()
+    ) as client:
+        await client.query("What's 123 * 456?")
+        async for _ in client.receive_response():
+            pass
 
-            await client.query("What time is it now?")
-            async for _ in client.receive_response():
-                pass
-
-    asyncio.run(_collect_messages())
+        await client.query("What time is it now?")
+        async for _ in client.receive_response():
+            pass
 
     spans_tuple = span_exporter.get_finished_spans()
     spans = sorted(list(spans_tuple), key=lambda x: x.start_time)