github-mirrors/tradingagents
1
0
Fork 0
mirror of https://github.com/TauricResearch/TradingAgents.git synced 2026-05-01 14:33:10 +03:00
Code Issues Packages Projects Releases Wiki Activity

refactor: simplify api_key mapping and consolidate tests

Browse Source 
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
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__":