From c05e10c39d8d1f1375decd2ac65d2c0f27842409 Mon Sep 17 00:00:00 2001
From: Anas Khan <83116240+anxkhn@users.noreply.github.com>
Date: Fri, 3 Jul 2026 16:25:31 +0530
Subject: [PATCH] test(models): exercise real connect() path in test_connect

test_connect patched google_llm.Gemini.connect (the method under test)
at the class level, so the real connect() body never ran and the only
assertion checked the test's own mock return value, a tautology that
could not fail on any regression.

Patch the _live_api_client boundary instead and let the real connect()
run, matching how the neighboring connect tests
(test_connect_with_custom_headers, test_connect_without_custom_headers,
test_connect_forwards_thinking_config) already work. Assert the live
client is called with the request model and live_connect_config, and
that connect() yields a GeminiLlmConnection wrapping the session with the
expected api_backend and model_version.

Signed-off-by: Anas Khan <83116240+anxkhn@users.noreply.github.com>
---
 tests/unittests/models/test_google_llm.py | 35 ++++++++++++++---------
 1 file changed, 21 insertions(+), 14 deletions(-)

diff --git a/tests/unittests/models/test_google_llm.py b/tests/unittests/models/test_google_llm.py
index caad24ed3f3..e63265153fc 100644
--- a/tests/unittests/models/test_google_llm.py
+++ b/tests/unittests/models/test_google_llm.py
@@ -558,25 +558,32 @@ async def test_generate_content_async_other_client_error(
 
 @pytest.mark.asyncio
 async def test_connect(gemini_llm, llm_request):
-  # Create a mock connection
-  mock_connection = mock.MagicMock(spec=GeminiLlmConnection)
+  """Test that connect yields a GeminiLlmConnection wrapping the live session."""
+  mock_live_session = mock.AsyncMock()
 
-  # Create a mock context manager
-  class MockContextManager:
+  # Patch the live API client boundary so the real connect() body runs.
+  with mock.patch.object(gemini_llm, "_live_api_client") as mock_live_client:
 
-    async def __aenter__(self):
-      return mock_connection
+    class MockLiveConnect:
 
-    async def __aexit__(self, *args):
-      pass
+      async def __aenter__(self):
+        return mock_live_session
+
+      async def __aexit__(self, *args):
+        pass
+
+    mock_live_client.aio.live.connect.return_value = MockLiveConnect()
 
-  # Mock the connect method at the class level
-  with mock.patch(
-      "google.adk.models.google_llm.Gemini.connect",
-      return_value=MockContextManager(),
-  ):
     async with gemini_llm.connect(llm_request) as connection:
-      assert connection is mock_connection
+      mock_live_client.aio.live.connect.assert_called_once()
+      call_args = mock_live_client.aio.live.connect.call_args
+      assert call_args.kwargs["model"] == llm_request.model
+      assert call_args.kwargs["config"] is llm_request.live_connect_config
+
+      assert isinstance(connection, GeminiLlmConnection)
+      assert connection._gemini_session is mock_live_session
+      assert connection._api_backend == gemini_llm._api_backend
+      assert connection._model_version == llm_request.model
 
 
 @pytest.mark.asyncio