TradingAgents/tests/conftest.py

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

import os
import tempfile
from unittest.mock import Mock

import pytest

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)