AI-Trader/tests/unit/test_calculate_final_position.py

"""Test _calculate_final_position_from_actions method."""

import pytest
from unittest.mock import patch
from agent.base_agent.base_agent import BaseAgent
from api.database import Database


@pytest.fixture
def test_db():
    """Create test database with schema."""
    db = Database(":memory:")

    # Create jobs record
    db.connection.execute("""
        INSERT INTO jobs (job_id, config_path, status, date_range, models, created_at)
        VALUES ('test-job', 'test.json', 'running', '2025-10-07 to 2025-10-07', 'gpt-5', '2025-10-07T00:00:00Z')
    """)
    db.connection.commit()

    return db


def test_calculate_final_position_first_day_with_trades(test_db):
    """Test calculating final position on first trading day with multiple trades."""

    # Create trading_day for first day
    trading_day_id = test_db.create_trading_day(
        job_id='test-job',
        model='gpt-5',
        date='2025-10-07',
        starting_cash=10000.0,
        starting_portfolio_value=10000.0,
        daily_profit=0.0,
        daily_return_pct=0.0,
        ending_cash=10000.0,  # Not yet calculated
        ending_portfolio_value=10000.0,  # Not yet calculated
        days_since_last_trading=1
    )

    # Add 15 buy actions (matching your real data)
    actions_data = [
        ("MSFT", 3, 528.285, "buy"),
        ("GOOGL", 6, 248.27, "buy"),
        ("NVDA", 10, 186.23, "buy"),
        ("LRCX", 6, 149.23, "buy"),
        ("AVGO", 2, 337.025, "buy"),
        ("AMZN", 5, 220.88, "buy"),
        ("MSFT", 2, 528.285, "buy"),  # Additional MSFT
        ("AMD", 4, 214.85, "buy"),
        ("CRWD", 1, 497.0, "buy"),
        ("QCOM", 4, 169.9, "buy"),
        ("META", 1, 717.72, "buy"),
        ("NVDA", 20, 186.23, "buy"),  # Additional NVDA
        ("NVDA", 13, 186.23, "buy"),  # Additional NVDA
        ("NVDA", 20, 186.23, "buy"),  # Additional NVDA
        ("NVDA", 53, 186.23, "buy"),  # Additional NVDA
    ]

    for symbol, quantity, price, action_type in actions_data:
        test_db.create_action(
            trading_day_id=trading_day_id,
            action_type=action_type,
            symbol=symbol,
            quantity=quantity,
            price=price
        )

    test_db.connection.commit()

    # Create BaseAgent instance
    agent = BaseAgent(signature="gpt-5", basemodel="anthropic/claude-sonnet-4", stock_symbols=[])

    # Mock Database() to return our test_db
    with patch('api.database.Database', return_value=test_db):
        # Calculate final position
        holdings, cash = agent._calculate_final_position_from_actions(
            trading_day_id=trading_day_id,
            starting_cash=10000.0
        )

    # Verify holdings
    assert holdings["MSFT"] == 5, f"Expected 5 MSFT (3+2) but got {holdings.get('MSFT', 0)}"
    assert holdings["GOOGL"] == 6, f"Expected 6 GOOGL but got {holdings.get('GOOGL', 0)}"
    assert holdings["NVDA"] == 116, f"Expected 116 NVDA (10+20+13+20+53) but got {holdings.get('NVDA', 0)}"
    assert holdings["LRCX"] == 6, f"Expected 6 LRCX but got {holdings.get('LRCX', 0)}"
    assert holdings["AVGO"] == 2, f"Expected 2 AVGO but got {holdings.get('AVGO', 0)}"
    assert holdings["AMZN"] == 5, f"Expected 5 AMZN but got {holdings.get('AMZN', 0)}"
    assert holdings["AMD"] == 4, f"Expected 4 AMD but got {holdings.get('AMD', 0)}"
    assert holdings["CRWD"] == 1, f"Expected 1 CRWD but got {holdings.get('CRWD', 0)}"
    assert holdings["QCOM"] == 4, f"Expected 4 QCOM but got {holdings.get('QCOM', 0)}"
    assert holdings["META"] == 1, f"Expected 1 META but got {holdings.get('META', 0)}"

    # Verify cash (should be less than starting)
    assert cash < 10000.0, f"Cash should be less than $10,000 but got ${cash}"

    # Calculate expected cash
    total_spent = sum(qty * price for _, qty, price, _ in actions_data)
    expected_cash = 10000.0 - total_spent
    assert abs(cash - expected_cash) < 0.01, f"Expected cash ${expected_cash} but got ${cash}"


def test_calculate_final_position_with_previous_holdings(test_db):
    """Test calculating final position when starting with existing holdings."""

    # Create day 1 with ending holdings
    day1_id = test_db.create_trading_day(
        job_id='test-job',
        model='gpt-5',
        date='2025-10-06',
        starting_cash=10000.0,
        starting_portfolio_value=10000.0,
        daily_profit=0.0,
        daily_return_pct=0.0,
        ending_cash=8000.0,
        ending_portfolio_value=9500.0,
        days_since_last_trading=1
    )

    # Add day 1 ending holdings
    test_db.create_holding(day1_id, "AAPL", 10)
    test_db.create_holding(day1_id, "MSFT", 5)

    # Create day 2
    day2_id = test_db.create_trading_day(
        job_id='test-job',
        model='gpt-5',
        date='2025-10-07',
        starting_cash=8000.0,
        starting_portfolio_value=9500.0,
        daily_profit=0.0,
        daily_return_pct=0.0,
        ending_cash=8000.0,
        ending_portfolio_value=9500.0,
        days_since_last_trading=1
    )

    # Add day 2 actions (buy more AAPL, sell some MSFT)
    test_db.create_action(day2_id, "buy", "AAPL", 5, 150.0)
    test_db.create_action(day2_id, "sell", "MSFT", 2, 500.0)

    test_db.connection.commit()

    # Create BaseAgent instance
    agent = BaseAgent(signature="gpt-5", basemodel="anthropic/claude-sonnet-4", stock_symbols=[])

    # Mock Database() to return our test_db
    with patch('api.database.Database', return_value=test_db):
        # Calculate final position for day 2
        holdings, cash = agent._calculate_final_position_from_actions(
            trading_day_id=day2_id,
            starting_cash=8000.0
        )

    # Verify holdings
    assert holdings["AAPL"] == 15, f"Expected 15 AAPL (10+5) but got {holdings.get('AAPL', 0)}"
    assert holdings["MSFT"] == 3, f"Expected 3 MSFT (5-2) but got {holdings.get('MSFT', 0)}"

    # Verify cash
    # Started: 8000
    # Buy 5 AAPL @ 150 = -750
    # Sell 2 MSFT @ 500 = +1000
    # Final: 8000 - 750 + 1000 = 8250
    expected_cash = 8000.0 - (5 * 150.0) + (2 * 500.0)
    assert abs(cash - expected_cash) < 0.01, f"Expected cash ${expected_cash} but got ${cash}"


def test_calculate_final_position_no_trades(test_db):
    """Test calculating final position when no trades were executed."""

    # Create day 1 with ending holdings
    day1_id = test_db.create_trading_day(
        job_id='test-job',
        model='gpt-5',
        date='2025-10-06',
        starting_cash=10000.0,
        starting_portfolio_value=10000.0,
        daily_profit=0.0,
        daily_return_pct=0.0,
        ending_cash=9000.0,
        ending_portfolio_value=10000.0,
        days_since_last_trading=1
    )

    test_db.create_holding(day1_id, "AAPL", 10)

    # Create day 2 with NO actions
    day2_id = test_db.create_trading_day(
        job_id='test-job',
        model='gpt-5',
        date='2025-10-07',
        starting_cash=9000.0,
        starting_portfolio_value=10000.0,
        daily_profit=0.0,
        daily_return_pct=0.0,
        ending_cash=9000.0,
        ending_portfolio_value=10000.0,
        days_since_last_trading=1
    )

    # No actions added
    test_db.connection.commit()

    # Create BaseAgent instance
    agent = BaseAgent(signature="gpt-5", basemodel="anthropic/claude-sonnet-4", stock_symbols=[])

    # Mock Database() to return our test_db
    with patch('api.database.Database', return_value=test_db):
        # Calculate final position
        holdings, cash = agent._calculate_final_position_from_actions(
            trading_day_id=day2_id,
            starting_cash=9000.0
        )

    # Verify holdings unchanged
    assert holdings["AAPL"] == 10, f"Expected 10 AAPL but got {holdings.get('AAPL', 0)}"

    # Verify cash unchanged
    assert abs(cash - 9000.0) < 0.01, f"Expected cash $9000 but got ${cash}"