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),)