TradingAgents/tests/conftest.py

"""Pytest configuration and shared fixtures for TradingAgents tests."""

import os
import pytest
import tempfile
from unittest.mock import Mock, MagicMock
from datetime import date, datetime
from typing import Dict, Any

from tradingagents.default_config import DEFAULT_CONFIG


@pytest.fixture
def sample_config():
    """Provide a test configuration."""
    config = DEFAULT_CONFIG.copy()
    config.update(
        {
            "online_tools": False,  # Use offline tools for testing
            "max_debate_rounds": 1,  # Limit rounds for faster tests
            "llm_provider": "openai",
            "deep_think_llm": "gpt-4o-mini",
            "quick_think_llm": "gpt-4o-mini",
            "project_dir": "/tmp/test_tradingagents",
        }
    )
    return config


@pytest.fixture
def mock_llm():
    """Mock LLM for testing."""
    mock = Mock()
    mock.model_name = "test-model"
    mock.invoke.return_value = Mock(content="Test response")
    mock.bind_tools.return_value = mock
    return mock


@pytest.fixture
def sample_ticker():
    """Sample stock ticker for testing."""
    return "AAPL"


@pytest.fixture
def sample_trade_date():
    """Sample trade date for testing."""
    return "2024-05-10"


@pytest.fixture
def sample_agent_state():
    """Sample agent state for testing."""
    return {
        "company_of_interest": "AAPL",
        "trade_date": "2024-05-10",
        "messages": [],
        "market_report": "",
        "sentiment_report": "",
        "news_report": "",
        "fundamentals_report": "",
        "investment_debate_state": {
            "bull_history": [],
            "bear_history": [],
            "history": [],
            "current_response": "",
            "judge_decision": "",
        },
        "trader_investment_plan": "",
        "risk_debate_state": {
            "risky_history": [],
            "safe_history": [],
            "neutral_history": [],
            "history": [],
            "judge_decision": "",
        },
        "investment_plan": "",
        "final_trade_decision": "",
    }


@pytest.fixture
def temp_data_dir():
    """Create a temporary directory for test data."""
    with tempfile.TemporaryDirectory() as temp_dir:
        yield temp_dir


@pytest.fixture
def mock_toolkit():
    """Mock toolkit with all necessary methods."""
    toolkit = Mock()
    toolkit.config = {"online_tools": False}

    # Mock data retrieval methods
    toolkit.get_YFin_data = Mock()
    toolkit.get_YFin_data_online = Mock()
    toolkit.get_stockstats_indicators_report = Mock()
    toolkit.get_stockstats_indicators_report_online = Mock()
    toolkit.get_reddit_stock_info = Mock()
    toolkit.get_stock_news_openai = Mock()
    toolkit.get_finnhub_news = Mock()
    toolkit.get_reddit_news = Mock()
    toolkit.get_global_news_openai = Mock()
    toolkit.get_google_news = Mock()
    toolkit.get_fundamentals_openai = Mock()
    toolkit.get_finnhub_company_insider_sentiment = Mock()
    toolkit.get_finnhub_company_insider_transactions = Mock()
    toolkit.get_simfin_balance_sheet = Mock()
    toolkit.get_simfin_cashflow = Mock()
    toolkit.get_simfin_income_stmt = Mock()

    return toolkit


@pytest.fixture
def sample_market_data():
    """Sample market data for testing."""
    return {
        "close": 150.0,
        "high": 155.0,
        "low": 148.0,
        "volume": 1000000,
        "date": "2024-05-10",
    }


@pytest.fixture
def sample_financial_data():
    """Sample financial data for testing."""
    return {
        "revenue": 100000000,
        "net_income": 20000000,
        "assets": 500000000,
        "liabilities": 200000000,
        "period": "annual",
        "date": "2024-05-10",
    }


@pytest.fixture(autouse=True)
def setup_test_environment(monkeypatch, temp_data_dir):
    """Set up test environment variables and directories."""
    # Set test environment variables
    monkeypatch.setenv("OPENAI_API_KEY", "test-key")
    monkeypatch.setenv("FINNHUB_API_KEY", "test-key")
    monkeypatch.setenv("REDDIT_CLIENT_ID", "test-id")
    monkeypatch.setenv("REDDIT_CLIENT_SECRET", "test-secret")

    # Create test data directories
    data_cache_dir = os.path.join(temp_data_dir, "dataflows", "data_cache")
    os.makedirs(data_cache_dir, exist_ok=True)

    finnhub_dir = os.path.join(data_cache_dir, "finnhub_data")
    os.makedirs(finnhub_dir, exist_ok=True)

    # Create subdirectories for different data types
    for data_type in ["news_data", "insider_trans", "SEC_filings", "fin_as_reported"]:
        os.makedirs(os.path.join(finnhub_dir, data_type), exist_ok=True)


@pytest.fixture
def mock_memory():
    """Mock financial situation memory."""
    memory = Mock()
    memory.add_memory = Mock()
    memory.get_memory = Mock(return_value="")
    memory.clear_memory = Mock()
    return memory


# Pytest configuration
def pytest_configure(config):
    """Configure pytest with custom markers."""
    config.addinivalue_line(
        "markers", "integration: mark test as integration test (slow)"
    )
    config.addinivalue_line("markers", "unit: mark test as unit test (fast)")
    config.addinivalue_line("markers", "api: mark test as requiring API access")


def pytest_collection_modifyitems(config, items):
    """Modify test collection to add markers based on location."""
    for item in items:
        # Add unit marker to tests in unit/ directory
        if "unit" in str(item.fspath):
            item.add_marker(pytest.mark.unit)
        # Add integration marker to tests in integration/ directory
        elif "integration" in str(item.fspath):
            item.add_marker(pytest.mark.integration)