From 047b38971cba6b390d04bb73e7191f2c05ee135e Mon Sep 17 00:00:00 2001
From: javierdejesusda <javier.dejesusj9@gmail.com>
Date: Tue, 24 Mar 2026 14:52:51 +0100
Subject: [PATCH] refactor: simplify api_key mapping and consolidate tests

Apply review suggestions: use concise `or` pattern for API key
resolution, consolidate tests into parameterized subTest, move
import to module level per PEP 8.
---
 tests/test_google_api_key.py               | 41 ++++++++--------------
 tradingagents/llm_clients/google_client.py |  7 ++--
 2 files changed, 18 insertions(+), 30 deletions(-)

diff --git a/tests/test_google_api_key.py b/tests/test_google_api_key.py
index 1ec2301f..e1607c49 100644
--- a/tests/test_google_api_key.py
+++ b/tests/test_google_api_key.py
@@ -1,38 +1,27 @@
 import unittest
 from unittest.mock import patch
 
+from tradingagents.llm_clients.google_client import GoogleClient
+
 
 class TestGoogleApiKeyStandardization(unittest.TestCase):
     """Verify GoogleClient accepts unified api_key parameter."""
 
     @patch("tradingagents.llm_clients.google_client.NormalizedChatGoogleGenerativeAI")
-    def test_api_key_mapped_to_google_api_key(self, mock_chat):
-        from tradingagents.llm_clients.google_client import GoogleClient
+    def test_api_key_handling(self, mock_chat):
+        test_cases = [
+            ("unified api_key is mapped", {"api_key": "test-key-123"}, "test-key-123"),
+            ("legacy google_api_key still works", {"google_api_key": "legacy-key-456"}, "legacy-key-456"),
+            ("unified api_key takes precedence", {"api_key": "unified", "google_api_key": "legacy"}, "unified"),
+        ]
 
-        client = GoogleClient("gemini-2.5-flash", api_key="test-key-123")
-        client.get_llm()
-        call_kwargs = mock_chat.call_args[1]
-        self.assertEqual(call_kwargs["google_api_key"], "test-key-123")
-
-    @patch("tradingagents.llm_clients.google_client.NormalizedChatGoogleGenerativeAI")
-    def test_legacy_google_api_key_still_works(self, mock_chat):
-        from tradingagents.llm_clients.google_client import GoogleClient
-
-        client = GoogleClient("gemini-2.5-flash", google_api_key="legacy-key-456")
-        client.get_llm()
-        call_kwargs = mock_chat.call_args[1]
-        self.assertEqual(call_kwargs["google_api_key"], "legacy-key-456")
-
-    @patch("tradingagents.llm_clients.google_client.NormalizedChatGoogleGenerativeAI")
-    def test_api_key_takes_precedence_over_google_api_key(self, mock_chat):
-        from tradingagents.llm_clients.google_client import GoogleClient
-
-        client = GoogleClient(
-            "gemini-2.5-flash", api_key="unified", google_api_key="legacy"
-        )
-        client.get_llm()
-        call_kwargs = mock_chat.call_args[1]
-        self.assertEqual(call_kwargs["google_api_key"], "unified")
+        for msg, kwargs, expected_key in test_cases:
+            with self.subTest(msg=msg):
+                mock_chat.reset_mock()
+                client = GoogleClient("gemini-2.5-flash", **kwargs)
+                client.get_llm()
+                call_kwargs = mock_chat.call_args[1]
+                self.assertEqual(call_kwargs.get("google_api_key"), expected_key)
 
 
 if __name__ == "__main__":
diff --git a/tradingagents/llm_clients/google_client.py b/tradingagents/llm_clients/google_client.py
index af8c6e48..f9971aa6 100644
--- a/tradingagents/llm_clients/google_client.py
+++ b/tradingagents/llm_clients/google_client.py
@@ -32,10 +32,9 @@ class GoogleClient(BaseLLMClient):
                 llm_kwargs[key] = self.kwargs[key]
 
         # Unified api_key maps to provider-specific google_api_key
-        if "api_key" in self.kwargs:
-            llm_kwargs["google_api_key"] = self.kwargs["api_key"]
-        elif "google_api_key" in self.kwargs:
-            llm_kwargs["google_api_key"] = self.kwargs["google_api_key"]
+        google_api_key = self.kwargs.get("api_key") or self.kwargs.get("google_api_key")
+        if google_api_key:
+            llm_kwargs["google_api_key"] = google_api_key
 
         # Map thinking_level to appropriate API param based on model
         # Gemini 3 Pro: low, high