fix: fix methods signature to pass updated tests

Author: yanomaly

src/strands/models/writer.py

@@ -7,7 +7,7 @@
 import json
 import logging
 import mimetypes
-from typing import Any, Dict, Generator, Iterable, List, Optional, Type, TypedDict, TypeVar, Union, cast
+from typing import Any, AsyncGenerator, Dict, List, Optional, Type, TypedDict, TypeVar, Union, cast
 
 import writerai
 from pydantic import BaseModel
@@ -349,7 +349,7 @@ def format_chunk(self, event: Any) -> StreamEvent:
                 raise RuntimeError(f"chunk_type=<{event['chunk_type']} | unknown type")
 
     @override
-    def stream(self, request: Any) -> Iterable[Any]:
+    async def stream(self, request: Any) -> AsyncGenerator[Any, None]:
         """Send the request to the model and get a streaming response.
 
         Args:
@@ -405,9 +405,9 @@ def stream(self, request: Any) -> Iterable[Any]:
         yield {"chunk_type": "metadata", "data": chunk.usage}
 
     @override
-    def structured_output(
+    async def structured_output(
         self, output_model: Type[T], prompt: Messages
-    ) -> Generator[dict[str, Union[T, Any]], None, None]:
+    ) -> AsyncGenerator[dict[str, Union[T, Any]], None]:
         """Get structured output from the model.
 
         Args:

tests/strands/models/test_writer.py

@@ -264,7 +264,8 @@ def test_format_request_with_unsupported_type(model, content, content_type):
         model.format_request(messages)
 
 
-def test_stream(writer_client, model, model_id):
+@pytest.mark.asyncio
+async def test_stream(writer_client, model, model_id):
     mock_tool_call_1_part_1 = unittest.mock.Mock(index=0)
     mock_tool_call_2_part_1 = unittest.mock.Mock(index=1)
     mock_delta_1 = unittest.mock.Mock(
@@ -292,7 +293,7 @@ def test_stream(writer_client, model, model_id):
     }
     response = model.stream(request)
 
-    events = list(response)
+    events = [event async for event in response]
     exp_events = [
         {"chunk_type": "message_start"},
         {"chunk_type": "content_block_start", "data_type": "text"},
@@ -313,7 +314,8 @@ def test_stream(writer_client, model, model_id):
     writer_client.chat.chat(**request)
 
 
-def test_stream_empty(writer_client, model, model_id):
+@pytest.mark.asyncio
+async def test_stream_empty(writer_client, model, model_id):
     mock_delta = unittest.mock.Mock(content=None, tool_calls=None)
     mock_usage = unittest.mock.Mock(prompt_tokens=0, completion_tokens=0, total_tokens=0)
 
@@ -327,7 +329,7 @@ def test_stream_empty(writer_client, model, model_id):
     request = {"model": model_id, "messages": [{"role": "user", "content": []}]}
     response = model.stream(request)
 
-    events = list(response)
+    events = [event async for event in response]
     exp_events = [
         {"chunk_type": "message_start"},
         {"chunk_type": "content_block_start", "data_type": "text"},
@@ -340,7 +342,8 @@ def test_stream_empty(writer_client, model, model_id):
     writer_client.chat.chat.assert_called_once_with(**request)
 
 
-def test_stream_with_empty_choices(writer_client, model, model_id):
+@pytest.mark.asyncio
+async def test_stream_with_empty_choices(writer_client, model, model_id):
     mock_delta = unittest.mock.Mock(content="content", tool_calls=None)
     mock_usage = unittest.mock.Mock(prompt_tokens=10, completion_tokens=20, total_tokens=30)
 
@@ -355,7 +358,7 @@ def test_stream_with_empty_choices(writer_client, model, model_id):
     request = {"model": model_id, "messages": [{"role": "user", "content": ["test"]}]}
     response = model.stream(request)
 
-    events = list(response)
+    events = [event async for event in response]
     exp_events = [
         {"chunk_type": "message_start"},
         {"chunk_type": "content_block_start", "data_type": "text"},