charieswei
/
TradingAgents
mirror of https://github.com/TauricResearch/TradingAgents.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

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.
This commit is contained in:
javierdejesusda 2026-03-24 14:52:51 +01:00
parent f5026009f9
commit 047b38971c
2 changed files with 18 additions and 30 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

41
tests/test_google_api_key.py
View File

@ -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__":

7
tradingagents/llm_clients/google_client.py
View File

@ -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