From cd43fdb48a30b043f91837f93391b9f8949f6fd4 Mon Sep 17 00:00:00 2001 From: coder-aditi Date: Wed, 2 Jul 2025 21:41:27 +0530 Subject: [PATCH 1/2] Fix: test_base_llm_flow.py - patch LLM flow streaming unit test --- .../adk/flows/llm_flows/base_llm_flow.py | 2 +- .../flows/llm_flows/test_base_llm_flow.py | 44 +++++++++++++++++++ 2 files changed, 45 insertions(+), 1 deletion(-) create mode 100644 tests/unittests/flows/llm_flows/test_base_llm_flow.py diff --git a/src/google/adk/flows/llm_flows/base_llm_flow.py b/src/google/adk/flows/llm_flows/base_llm_flow.py index be1a1571a..8bcdb7de7 100644 --- a/src/google/adk/flows/llm_flows/base_llm_flow.py +++ b/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 diff --git a/tests/unittests/flows/llm_flows/test_base_llm_flow.py b/tests/unittests/flows/llm_flows/test_base_llm_flow.py new file mode 100644 index 000000000..35b510791 --- /dev/null +++ b/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" From e1a1a9ddb6deaf8a367e9eabd3b23797816e0417 Mon Sep 17 00:00:00 2001 From: coder-aditi Date: Wed, 2 Jul 2025 21:50:59 +0530 Subject: [PATCH 2/2] Trigger CLA recheck