aitrader/tests/test_portfolio.py

import pytest
from unittest.mock import MagicMock

from aitrader.config import Settings
from aitrader.storage import db as dbm
from aitrader.storage.models import Trade
from aitrader.trader import portfolio


@pytest.fixture
def settings(tmp_path):
    s = Settings(starting_equity_eur=10000.0, db_path=str(tmp_path / "t.db"))
    # Initialize DB
    dbm.get_engine(s.db_path)
    return s


@pytest.fixture
def mock_kraken():
    kraken = MagicMock()
    kraken.create_market_order.return_value = {"id": "mock_order_123"}
    return kraken


def test_check_stop_take_profit_buy_sl_hit(settings, mock_kraken):
    # Setup open buy trade
    with dbm.session(settings.db_path) as s:
        trade = Trade(
            symbol="BTC/USD:USD",
            side="buy",
            qty=1.5,
            entry_price=60000.0,
            stop_loss=59000.0,
            take_profit=62000.0,
            status="open"
        )
        s.add(trade)
        s.commit()
        s.refresh(trade)
        trade_id = trade.id

    # Under SL (58500 <= 59000)
    closed = portfolio.check_stop_take_profit(settings, mock_kraken, "BTC/USD:USD", 58500.0)
    
    assert closed == [trade_id]
    mock_kraken.create_market_order.assert_called_once_with("BTC/USD:USD", "sell", 1.5)

    # Check DB update
    with dbm.session(settings.db_path) as s:
        t = s.get(Trade, trade_id)
        assert t.status == "closed"
        assert t.exit_price == 58500.0
        assert t.pnl_eur == (58500.0 - 60000.0) * 1.5


def test_check_stop_take_profit_buy_tp_hit(settings, mock_kraken):
    with dbm.session(settings.db_path) as s:
        trade = Trade(
            symbol="BTC/USD:USD",
            side="buy",
            qty=1.0,
            entry_price=60000.0,
            stop_loss=59000.0,
            take_profit=62000.0,
            status="open"
        )
        s.add(trade)
        s.commit()
        s.refresh(trade)
        trade_id = trade.id

    # Above TP (62500 >= 62000)
    closed = portfolio.check_stop_take_profit(settings, mock_kraken, "BTC/USD:USD", 62500.0)
    
    assert closed == [trade_id]
    mock_kraken.create_market_order.assert_called_once_with("BTC/USD:USD", "sell", 1.0)

    # Check DB update
    with dbm.session(settings.db_path) as s:
        t = s.get(Trade, trade_id)
        assert t.status == "closed"
        assert t.exit_price == 62500.0
        assert t.pnl_eur == (62500.0 - 60000.0) * 1.0


def test_check_stop_take_profit_no_hit(settings, mock_kraken):
    with dbm.session(settings.db_path) as s:
        trade = Trade(
            symbol="BTC/USD:USD",
            side="buy",
            qty=1.0,
            entry_price=60000.0,
            stop_loss=59000.0,
            take_profit=62000.0,
            status="open"
        )
        s.add(trade)
        s.commit()
        s.refresh(trade)
        trade_id = trade.id

    # Price between SL and TP
    closed = portfolio.check_stop_take_profit(settings, mock_kraken, "BTC/USD:USD", 60500.0)
    
    assert closed == []
    mock_kraken.create_market_order.assert_not_called()

    # Check DB update
    with dbm.session(settings.db_path) as s:
        t = s.get(Trade, trade_id)
        assert t.status == "open"


def test_check_stop_take_profit_sell_sl_hit(settings, mock_kraken):
    # Setup open sell trade
    with dbm.session(settings.db_path) as s:
        trade = Trade(
            symbol="ETH/USD:USD",
            side="sell",
            qty=2.0,
            entry_price=3000.0,
            stop_loss=3100.0,
            take_profit=2800.0,
            status="open"
        )
        s.add(trade)
        s.commit()
        s.refresh(trade)
        trade_id = trade.id

    # Above SL for Sell (3150 >= 3100)
    closed = portfolio.check_stop_take_profit(settings, mock_kraken, "ETH/USD:USD", 3150.0)
    
    assert closed == [trade_id]
    mock_kraken.create_market_order.assert_called_once_with("ETH/USD:USD", "buy", 2.0)

    # Check DB update
    with dbm.session(settings.db_path) as s:
        t = s.get(Trade, trade_id)
        assert t.status == "closed"
        assert t.exit_price == 3150.0
        assert t.pnl_eur == -1 * (3150.0 - 3000.0) * 2.0


def test_check_stop_take_profit_exchange_fails(settings, mock_kraken):
    mock_kraken.create_market_order.side_effect = Exception("Kraken API down")
    
    with dbm.session(settings.db_path) as s:
        trade = Trade(
            symbol="BTC/USD:USD",
            side="buy",
            qty=1.0,
            entry_price=60000.0,
            stop_loss=59000.0,
            take_profit=62000.0,
            status="open"
        )
        s.add(trade)
        s.commit()
        s.refresh(trade)
        trade_id = trade.id

    # Under SL, but exchange order fails
    closed = portfolio.check_stop_take_profit(settings, mock_kraken, "BTC/USD:USD", 58500.0)
    
    assert closed == []
    mock_kraken.create_market_order.assert_called_once_with("BTC/USD:USD", "sell", 1.0)

    # DB record must still be open
    with dbm.session(settings.db_path) as s:
        t = s.get(Trade, trade_id)
        assert t.status == "open"