"""Tests for Portfolio Manager reading the correct state key. The Portfolio Manager must use state["trader_investment_plan"] (the Trader's memory-refined output), not state["investment_plan"] (the Research Manager's raw output). Using the wrong key means the Trader agent's entire contribution — including lessons from past trades — is silently bypassed. """ import types import unittest from unittest.mock import MagicMock class TestPortfolioManagerStateKey(unittest.TestCase): """Verify Portfolio Manager reads trader_investment_plan, not investment_plan.""" def _make_state(self, trader_plan="TRADER PLAN: BUY AAPL", research_plan="RESEARCH PLAN: HOLD AAPL"): """Build a minimal AgentState-like dict for testing.""" return { "company_of_interest": "AAPL", "trade_date": "2025-01-15", "market_report": "Market is bullish with strong momentum.", "sentiment_report": "Social sentiment is positive.", "news_report": "No significant negative news.", "fundamentals_report": "Strong earnings, low P/E ratio.", "investment_plan": research_plan, # Research Manager's output "trader_investment_plan": trader_plan, # Trader's refined output "risk_debate_state": { "history": "Aggressive: Go all in. Conservative: Be careful.", "aggressive_history": "Aggressive: Go all in.", "conservative_history": "Conservative: Be careful.", "neutral_history": "Neutral: Balance risk.", "latest_speaker": "Neutral", "current_aggressive_response": "Aggressive: Go all in.", "current_conservative_response": "Conservative: Be careful.", "current_neutral_response": "Neutral: Balance risk.", "judge_decision": "", "count": 3, }, } def test_portfolio_manager_uses_trader_plan_not_research_plan(self): """The prompt sent to the LLM must contain the Trader's plan, not the Research Manager's.""" from tradingagents.agents.managers.portfolio_manager import create_portfolio_manager mock_llm = MagicMock() mock_response = MagicMock() mock_response.content = "Rating: Buy. The trader's analysis is compelling." mock_llm.invoke.return_value = mock_response mock_memory = MagicMock() mock_memory.get_memories.return_value = [] node = create_portfolio_manager(mock_llm, mock_memory) state = self._make_state( trader_plan="UNIQUE_TRADER_MARKER: BUY with 60% position", research_plan="UNIQUE_RESEARCH_MARKER: HOLD with caution", ) node(state) # The LLM should have been called with the Trader's plan call_args = mock_llm.invoke.call_args prompt_text = call_args[0][0] if call_args[0] else str(call_args) self.assertIn("UNIQUE_TRADER_MARKER", prompt_text, "Portfolio Manager must use trader_investment_plan (Trader's output), " "not investment_plan (Research Manager's output)") self.assertNotIn("UNIQUE_RESEARCH_MARKER", prompt_text, "Portfolio Manager should NOT contain Research Manager's raw plan") def test_portfolio_manager_returns_final_trade_decision(self): """Portfolio Manager must return final_trade_decision in its output.""" from tradingagents.agents.managers.portfolio_manager import create_portfolio_manager mock_llm = MagicMock() mock_response = MagicMock() mock_response.content = "Final decision: BUY AAPL at market open." mock_llm.invoke.return_value = mock_response mock_memory = MagicMock() mock_memory.get_memories.return_value = [] node = create_portfolio_manager(mock_llm, mock_memory) state = self._make_state() result = node(state) self.assertIn("final_trade_decision", result) self.assertEqual(result["final_trade_decision"], "Final decision: BUY AAPL at market open.") def test_portfolio_manager_preserves_risk_debate_state(self): """Portfolio Manager must preserve all risk_debate_state fields.""" from tradingagents.agents.managers.portfolio_manager import create_portfolio_manager mock_llm = MagicMock() mock_response = MagicMock() mock_response.content = "Decision rendered." mock_llm.invoke.return_value = mock_response mock_memory = MagicMock() mock_memory.get_memories.return_value = [] node = create_portfolio_manager(mock_llm, mock_memory) state = self._make_state() result = node(state) rds = result["risk_debate_state"] self.assertEqual(rds["judge_decision"], "Decision rendered.") self.assertEqual(rds["latest_speaker"], "Judge") # Original histories must be preserved self.assertEqual(rds["aggressive_history"], "Aggressive: Go all in.") self.assertEqual(rds["conservative_history"], "Conservative: Be careful.") self.assertEqual(rds["neutral_history"], "Neutral: Balance risk.") self.assertEqual(rds["count"], 3) def test_distinct_trader_and_research_plans_both_present_in_state(self): """Ensure the test state correctly models both plan types as distinct values.""" state = self._make_state( trader_plan="Trader says BUY aggressively", research_plan="Research says HOLD cautiously", ) self.assertNotEqual(state["investment_plan"], state["trader_investment_plan"], "Test setup: plans must be distinct to catch the wrong-key bug") def test_portfolio_manager_with_memory_recommendations(self): """Past memory recommendations should appear in the prompt.""" from tradingagents.agents.managers.portfolio_manager import create_portfolio_manager mock_llm = MagicMock() mock_response = MagicMock() mock_response.content = "Decision with memory context." mock_llm.invoke.return_value = mock_response mock_memory = MagicMock() mock_memory.get_memories.return_value = [ {"recommendation": "PAST_LESSON_ALPHA: avoid overleveraging", "similarity_score": 0.9}, {"recommendation": "PAST_LESSON_BETA: check earnings date", "similarity_score": 0.7}, ] node = create_portfolio_manager(mock_llm, mock_memory) state = self._make_state() node(state) prompt_text = mock_llm.invoke.call_args[0][0] self.assertIn("PAST_LESSON_ALPHA", prompt_text) self.assertIn("PAST_LESSON_BETA", prompt_text) class TestStockStatsBulkImport(unittest.TestCase): """Verify _get_stock_stats_bulk has access to pandas.""" def test_pandas_imported_in_y_finance(self): """y_finance.py must import pandas for _get_stock_stats_bulk to work.""" import importlib import tradingagents.dataflows.y_finance as yf_module # pd must be accessible in the module namespace self.assertTrue( hasattr(yf_module, 'pd'), "_get_stock_stats_bulk uses pd.isna() but pandas is not imported as 'pd' " "in y_finance.py, causing NameError on every call" ) def test_pd_isna_callable(self): """pd.isna must be callable from y_finance module scope.""" import tradingagents.dataflows.y_finance as yf_module pd_ref = getattr(yf_module, 'pd', None) self.assertIsNotNone(pd_ref, "pd not found in y_finance module") self.assertTrue(callable(getattr(pd_ref, 'isna', None)), "pd.isna must be callable") if __name__ == "__main__": unittest.main()