trex-core/tests/application/test_reconciliation.py

"""Position reconciliation tests (REQ-031)."""
from __future__ import annotations

from unittest.mock import Mock

import pytest

from trex_core.application.broker_port import BalanceItem, BrokerBalanceError
from trex_core.application.portfolio import (
    InMemoryReconcileWarningStore,
    PortfolioItem,
    SnapshotStore,
)
from trex_core.application.reconciliation import reconcile_positions


class FakeBroker:
    def __init__(self, balances):
        self._b = balances
    def name(self): return "fake"
    def submit_order(self, req): ...
    def poll_fill(self, *args, **kw): ...
    def cancel_order(self, *args, **kw): ...
    def get_balance(self): return list(self._b)


def _store_with(*items):
    store = SnapshotStore()
    store.swap(list(items))
    return store


def test_mismatch_raises_warning_and_masks_asset():
    store = _store_with(
        PortfolioItem(portfolio_id="pf-A", symbol="005930",
                      cash_balance=0, asset_quantity=10),
    )
    broker = FakeBroker([BalanceItem(symbol="005930", hldg_qty=7)])  # mismatch
    decisions = reconcile_positions(store=store, broker=broker)
    assert decisions["pf-A"] == "warning"
    assert store.warnings() == {"pf-A"}


def test_subsequent_match_auto_clears_warning():
    store = _store_with(
        PortfolioItem(portfolio_id="pf-A", symbol="005930",
                      cash_balance=0, asset_quantity=10),
    )
    # round 1 — mismatch
    reconcile_positions(store=store, broker=FakeBroker([BalanceItem("005930", 7)]))
    assert "pf-A" in store.warnings()
    # round 2 — match
    decisions = reconcile_positions(store=store, broker=FakeBroker([BalanceItem("005930", 10)]))
    assert decisions["pf-A"] == "cleared"
    assert store.warnings() == set()


def test_reconcile_positions_accepts_symbol_from_complete_later_page_snapshot(
    kis_balance_broker_factory,
):
    store = _store_with(
        PortfolioItem(portfolio_id="pf-later", symbol="000660", cash_balance=0, asset_quantity=4),
    )
    broker, calls = kis_balance_broker_factory([
        {"rows": [], "tr_cont": "F", "fk": "A", "nk": "1"},
        {"rows": [{"pdno": "000660", "hldg_qty": "4"}], "tr_cont": "E"},
    ])

    decisions = reconcile_positions(store=store, broker=broker)

    assert decisions == {"pf-later": "ok"}
    assert store.warnings() == set()
    assert len(calls) == 2


def test_multiple_portfolios_same_symbol_aggregate():
    store = _store_with(
        PortfolioItem(portfolio_id="pf-A", symbol="005930",
                      cash_balance=0, asset_quantity=6),
        PortfolioItem(portfolio_id="pf-B", symbol="005930",
                      cash_balance=0, asset_quantity=4),
    )
    # hub total 6+4=10 == broker 10 → both ok
    decisions = reconcile_positions(
        store=store, broker=FakeBroker([BalanceItem("005930", 10)]),
    )
    assert decisions == {"pf-A": "ok", "pf-B": "ok"}
    assert store.warnings() == set()


def test_paper_broker_returns_deterministic_balance():
    """REQ-031 AC-3: paper stub provides deterministic balance."""
    from trex_core.adapters.broker.paper import PaperBroker
    from trex_core.application.broker_port import OrderRequest

    broker = PaperBroker(fill_after_polls=0)
    ack = broker.submit_order(OrderRequest(slot_id="s", side="buy",
                                           symbol="005930", qty=10, limit_price=1000))
    broker.poll_fill(ack.broker_order_id, side="buy", symbol="005930")  # fill
    balances = broker.get_balance()
    assert len(balances) == 1
    assert balances[0].symbol == "005930"
    assert balances[0].hldg_qty == 10

    # zero-qty rows dropped (D-2 0-qty leftover semantics)
    ack2 = broker.submit_order(OrderRequest(slot_id="s2", side="buy",
                                            symbol="000660", qty=5, limit_price=500))
    broker.poll_fill(ack2.broker_order_id, side="buy", symbol="000660")
    ack3 = broker.submit_order(OrderRequest(slot_id="s3", side="sell",
                                            symbol="000660", qty=5, limit_price=500))
    broker.poll_fill(ack3.broker_order_id, side="sell", symbol="000660")
    syms = {b.symbol for b in broker.get_balance()}
    assert "000660" not in syms        # net 0 → dropped
    assert "005930" in syms


def test_reconcile_positions_propagates_later_page_failure_without_warning_write(
    kis_balance_broker_factory, monkeypatch
):
    store = SnapshotStore()
    store.swap([
        PortfolioItem(portfolio_id="pf-A", symbol="005930", cash_balance=0, asset_quantity=10),
    ])
    store.set_warning("pf-A", True)
    warning_write = Mock(wraps=store.set_warning)
    monkeypatch.setattr(store, "set_warning", warning_write)
    warn_store = InMemoryReconcileWarningStore()
    set_warned = Mock(wraps=warn_store.set_warned)
    clear_warned = Mock(wraps=warn_store.clear_warned)
    monkeypatch.setattr(warn_store, "set_warned", set_warned)
    monkeypatch.setattr(warn_store, "clear_warned", clear_warned)
    broker, calls = kis_balance_broker_factory([
        {
            "rows": [{"pdno": "005930", "hldg_qty": "10"}],
            "tr_cont": "F",
            "fk": "A",
            "nk": "1",
        },
        {"rows": [], "tr_cont": "X"},
    ])

    with pytest.raises(BrokerBalanceError, match="malformed_response"):
        reconcile_positions(
            store=store,
            broker=broker,
            slot_repo=object(),
            warn_store=warn_store,
        )

    assert len(calls) == 2
    assert store.warnings() == {"pf-A"}
    warning_write.assert_not_called()
    set_warned.assert_not_called()
    clear_warned.assert_not_called()