Fix: test_base_llm_flow.py - patch LLM flow streaming unit test

Author: coder-aditi

src/google/adk/flows/llm_flows/base_llm_flow.py

@@ -217,7 +217,7 @@ async def _receive_from_model(
   ) -> AsyncGenerator[Event, None]:
     """Receive data from model and process events using BaseLlmConnection."""
 
-    def get_author_for_event(llm_response):
+    def get_author_for_event(llm_response: LlmResponse) -> str:
       """Get the author of the event.
 
       When the model returns transcription, the author is "user". Otherwise, the

tests/unittests/flows/llm_flows/test_base_llm_flow.py

@@ -0,0 +1,44 @@
+import pytest
+from unittest.mock import MagicMock
+from google.adk.flows.llm_flows.base_llm_flow import BaseLlmFlow, LlmResponse, Event, ConnectionClosedOK
+
+@pytest.mark.asyncio
+async def test_receive_from_model_yields_events():
+    flow = BaseLlmFlow()
+
+    fake_response_1 = LlmResponse(
+        content=MagicMock(role='assistant'),
+        error_code=None,
+        interrupted=False
+    )
+    fake_response_2 = LlmResponse(
+        content=MagicMock(role='user'),
+        error_code=None,
+        interrupted=False
+    )
+
+    async def fake_receive():
+        yield fake_response_1
+        yield fake_response_2
+        raise ConnectionClosedOK(rcvd=None, sent=None)
+
+
+    llm_connection = MagicMock()
+    llm_connection.receive = fake_receive
+
+    invocation_context = MagicMock()
+    invocation_context.agent.name = "TestAgent"
+    invocation_context.live_request_queue = MagicMock()
+    invocation_context.transcription_cache = []
+    invocation_context.invocation_id = "test_invocation_id_123"
+
+    events = []
+    async for event in flow._receive_from_model(
+        llm_connection, event_id="test_event", invocation_context=invocation_context, llm_request=MagicMock()
+    ):
+        events.append(event)
+
+    # Add your assertions here, e.g.:
+    assert len(events) == 2
+    assert events[0].author == "TestAgent"
+    assert events[1].author == "user"