trex-core/tests/adapters/broker/test_ownership.py

from decimal import Decimal
from datetime import datetime, timezone

import pytest

from trex_core.adapters.broker.ownership import (
    OwnershipAwareBroker,
    StrategyOwnershipStore,
)
from trex_core.application.broker_port import (
    BalanceItem,
    BrokerBalanceError,
    CancelOutcome,
    CancelState,
    FillStatus,
    OrderAck,
)
from trex_core.application.account_snapshot import (
    AccountSnapshotService,
    InMemoryAccountSnapshotStore,
)


class FakeBroker:
    def __init__(self, balance, fill=None):
        self.balance = balance
        self.fill = fill or FillStatus(state="accepted")
        self.poll_metadata = None

    def name(self):
        return "fake"

    def submit_order(self, req):
        return OrderAck(accepted=True, broker_order_id="1")

    def poll_fill(
        self,
        broker_order_id,
        *,
        side,
        symbol,
        order_date=None,
        submitted_at=None,
    ):
        self.poll_metadata = (order_date, submitted_at)
        return self.fill

    def cancel_order(self, broker_order_id, *, side, symbol, qty, venue="KRX"):
        return CancelOutcome(CancelState.REQUESTED)

    def get_balance(self):
        return self.balance

    def account_scope_hash(self):
        return "c" * 64

    def get_available_cash(self, *, price_check_code="005930"):
        return 1_000_000

    def get_current_price(self, symbol, market="J"):
        return Decimal("1")


def test_manual_balance_is_hidden_until_trex_fill(tmp_path):
    store = StrategyOwnershipStore(tmp_path / "ownership.json")
    broker = OwnershipAwareBroker(
        FakeBroker([BalanceItem("005930", 19), BalanceItem("010120", 5)]),
        store,
    )
    assert broker.get_balance() == []


def test_confirmed_fill_is_persisted_once_and_manual_remainder_stays_hidden(tmp_path):
    store = StrategyOwnershipStore(tmp_path / "ownership.json")
    fill = FillStatus(
        state="filled",
        filled_qty=2,
        execution_id="kis:1:2",
    )
    delegate = FakeBroker([BalanceItem("005930", 21)], fill)
    broker = OwnershipAwareBroker(delegate, store)

    broker.poll_fill("1", side="buy", symbol="005930")
    broker.poll_fill("1", side="buy", symbol="005930")

    assert broker.get_balance() == [BalanceItem("005930", 2)]
    assert StrategyOwnershipStore(store.path).snapshot()[0] == {"005930": 2}


def test_poll_preserves_original_order_metadata_through_ownership_decorator(tmp_path):
    store = StrategyOwnershipStore(tmp_path / "ownership.json")
    delegate = FakeBroker([])
    broker = OwnershipAwareBroker(delegate, store)
    submitted_at = datetime(2026, 8, 14, 6, 20, tzinfo=timezone.utc)

    broker.poll_fill(
        "1",
        side="buy",
        symbol="005930",
        order_date="20260814",
        submitted_at=submitted_at,
    )

    assert delegate.poll_metadata == ("20260814", submitted_at)


def test_strategy_quantity_above_raw_broker_total_fails_closed(tmp_path):
    store = StrategyOwnershipStore(tmp_path / "ownership.json")
    store.apply_fill(
        execution_id="buy-1",
        symbol="005930",
        side="buy",
        quantity=4,
    )
    broker = OwnershipAwareBroker(FakeBroker([BalanceItem("005930", 3)]), store)

    with pytest.raises(BrokerBalanceError) as exc:
        broker.get_balance()
    assert exc.value.code == "ownership_mismatch"


def _paged_kis_for_ownership(kis_balance_broker_factory, strategy_qty):
    return kis_balance_broker_factory([
        {
            "rows": [{"pdno": "005930", "hldg_qty": "19"}],
            "tr_cont": "F",
            "fk": "A",
            "nk": "1",
        },
        {
            "rows": [{"pdno": "000660", "hldg_qty": str(strategy_qty)}],
            "tr_cont": "E",
        },
    ])


def test_projection_includes_strategy_symbol_from_actual_later_kis_page(
    tmp_path, kis_balance_broker_factory
):
    store = StrategyOwnershipStore(tmp_path / "ownership.json")
    store.apply_fill(
        execution_id="buy-1",
        symbol="000660",
        side="buy",
        quantity=4,
    )
    delegate, calls = _paged_kis_for_ownership(kis_balance_broker_factory, 4)
    broker = OwnershipAwareBroker(delegate, store)

    assert broker.get_balance() == [BalanceItem("000660", 4)]
    assert len(calls) == 2
    assert "tr_cont" not in calls[0]["headers"]
    assert calls[1]["headers"]["tr_cont"] == "N"


def test_projection_detects_mismatch_from_actual_later_kis_page(
    tmp_path, kis_balance_broker_factory
):
    store = StrategyOwnershipStore(tmp_path / "ownership.json")
    store.apply_fill(
        execution_id="buy-1",
        symbol="000660",
        side="buy",
        quantity=4,
    )
    delegate, calls = _paged_kis_for_ownership(kis_balance_broker_factory, 3)
    broker = OwnershipAwareBroker(delegate, store)

    with pytest.raises(BrokerBalanceError) as exc:
        broker.get_balance()
    assert exc.value.code == "ownership_mismatch"
    assert len(calls) == 2


def test_sell_cannot_reduce_below_strategy_owned_zero(tmp_path):
    store = StrategyOwnershipStore(tmp_path / "ownership.json")
    with pytest.raises(BrokerBalanceError):
        store.apply_fill(
            execution_id="sell-1",
            symbol="005930",
            side="sell",
            quantity=1,
        )


def test_account_snapshot_keeps_manual_holdings_outside_strategy_scope(tmp_path):
    ownership = StrategyOwnershipStore(tmp_path / "ownership.json")
    ownership.apply_fill(
        execution_id="buy-1",
        symbol="005930",
        side="buy",
        quantity=2,
    )
    broker = OwnershipAwareBroker(
        FakeBroker([BalanceItem("005930", 21), BalanceItem("000660", 1)]),
        ownership,
    )
    snapshots = InMemoryAccountSnapshotStore()
    service = AccountSnapshotService(
        broker=broker,
        store=snapshots,
        clock=lambda: datetime(2026, 7, 29, tzinfo=timezone.utc),
    )

    assert service.refresh()
    snapshot = snapshots.load("c" * 64)
    assert snapshot is not None
    assert snapshot.holdings == (BalanceItem("005930", 2),)