trex-core/tests/application/test_execution_cancel.py

"""External-cancel + stale-watchdog execution tests (dcy-013 / issue #9).

A programmable broker lets each poll return a chosen FillStatus so we can drive
full external cancel (-> rollback to holding) and partial-fill-then-cancel
(-> ledger record + qty decrement + holding remainder), plus the Alert-only
pending stale watchdog.
"""
from __future__ import annotations

import logging
from dataclasses import replace
from datetime import datetime, timedelta, timezone

from trex_core.adapters.ledger.in_memory import InMemoryLedger
from trex_core.adapters.persistence.in_memory_intent_store import InMemoryIntentStore
from trex_core.application.broker_port import BalanceItem, FillStatus, OrderAck
from trex_core.application.execution import ExecutionService
from trex_core.application.intents import BuyIntent, IntentService, SellIntent
from trex_core.application.ledger_outbox import InMemoryLedgerOutbox, LedgerWriter
from trex_core.domain.slot import SlotService, SlotState
from trex_core.domain.slot.in_memory_repository import InMemorySlotRepository



class ProgrammableBroker:
    """Broker double whose poll_fill response is chosen per call by `responder`."""

    def __init__(self):
        self._seq = 0
        # default: full fill (used to drive a buy to holding)
        self.responder = lambda side: FillStatus(
            state="filled", filled_qty=10, avg_price=1000.0,
            execution_id="exec-full", fee=1.0, tax=0.0,
        )
        self.polls = 0
        self.balance = {"005930": 10}

    def name(self):
        return "stub"

    def submit_order(self, req):
        self._seq += 1
        return OrderAck(accepted=True, broker_order_id=f"o{self._seq}")

    def poll_fill(self, broker_order_id, *, side, symbol):
        self.polls += 1
        return self.responder(side)

    def cancel_order(self, broker_order_id, *, side, symbol, qty):
        return True

    def get_balance(self):
        # M2 (dcy-020/REQ-062): sell dispatch reads the broker balance.
        return [BalanceItem(symbol=s, hldg_qty=q) for s, q in self.balance.items() if q > 0]

    def get_current_price(self, symbol):
        return None


def _wire(broker, *, stale_sec=None):
    repo = InMemorySlotRepository()
    svc = SlotService(repo)
    intents = IntentService(svc, repo, InMemoryIntentStore(),
                            resolve_symbol=lambda pid: "005930")
    outbox = InMemoryLedgerOutbox()
    ledger = InMemoryLedger()
    execution = ExecutionService(
        svc, repo, broker, ledger_outbox=outbox,
        stale_pending_threshold_sec=stale_sec,
    )
    writer = LedgerWriter(outbox, ledger)
    repo.reconcile_universe("default", {"PF-1"})
    return repo, svc, intents, execution, ledger, writer


def _buy_to_holding(intents, execution):
    sid = intents.submit_buy(BuyIntent(
        intent_id="b1", qty=10, entry_price=1000, portfolio_id="PF-1",
    )).slot_id
    execution.run_once()  # buy dispatched + full fill -> holding
    return sid


def test_full_external_cancel_returns_slot_to_holding():
    broker = ProgrammableBroker()
    repo, svc, intents, execution, ledger, writer = _wire(broker)
    sid = _buy_to_holding(intents, execution)
    assert repo.fetch_slot(sid).status is SlotState.HOLDING

    intents.submit_sell(SellIntent(intent_id="s1", slot_id=sid))
    broker.responder = lambda side: FillStatus(state="cancelled")  # external cancel
    r = execution.run_once()  # dispatch sell + poll(cancelled) -> rollback
    slot = repo.fetch_slot(sid)
    assert slot.status is SlotState.HOLDING       # recovered, not stuck pending_sell
    assert slot.broker_order_id is None           # dispatch cleared -> re-sellable
    assert r.rolled_back == 1


def test_partial_fill_then_cancel_records_and_keeps_remainder():
    broker = ProgrammableBroker()
    repo, svc, intents, execution, ledger, writer = _wire(broker)
    sid = _buy_to_holding(intents, execution)
    writer.run_once()  # drain the buy ledger item

    intents.submit_sell(SellIntent(intent_id="s1", slot_id=sid))
    # 4 of 10 sold, then externally cancelled
    broker.responder = lambda side: FillStatus(
        state="partially_filled", filled_qty=4, avg_price=1100.0,
        execution_id="exec-partial", fee=0.5, tax=0.0, cancelled=True,
    )
    r = execution.run_once()
    slot = repo.fetch_slot(sid)
    assert slot.status is SlotState.HOLDING        # remainder stays held
    assert slot.qty == 6                            # 10 - 4 sold
    assert slot.avg_price == 1000.0                 # cost basis preserved (not sell price)
    assert slot.broker_order_id is None
    assert r.filled == 1 and r.ledger_enqueued == 1

    writer.run_once()
    sells = [t for t in ledger.trades() if t.side == "sell"]
    assert len(sells) == 1 and sells[0].qty == 4 and sells[0].price == 1100.0


def test_stale_pending_watchdog_warns_without_state_change(caplog):
    broker = ProgrammableBroker()
    repo, svc, intents, execution, ledger, writer = _wire(broker, stale_sec=60)
    sid = _buy_to_holding(intents, execution)
    intents.submit_sell(SellIntent(intent_id="s1", slot_id=sid))
    # order sits accepted (never fills)
    broker.responder = lambda side: FillStatus(state="accepted")
    execution.run_once()  # dispatch sell, poll -> accepted (fresh, not stale yet)
    # backdate the dispatched slot's updated_at beyond the threshold
    s = repo.fetch_slot(sid)
    repo._slots[sid] = replace(s, updated_at=datetime.now(tz=timezone.utc) - timedelta(seconds=120))

    with caplog.at_level(logging.WARNING):
        r = execution.run_once()
    assert r.stale_pending == 1
    assert repo.fetch_slot(sid).status is SlotState.PENDING_SELL  # state unchanged
    assert any("pending order stale" in rec.message for rec in caplog.records)


def test_stale_watchdog_disabled_by_default():
    broker = ProgrammableBroker()
    repo, svc, intents, execution, ledger, writer = _wire(broker)  # no threshold
    sid = _buy_to_holding(intents, execution)
    intents.submit_sell(SellIntent(intent_id="s1", slot_id=sid))
    broker.responder = lambda side: FillStatus(state="accepted")
    execution.run_once()
    s = repo.fetch_slot(sid)
    repo._slots[sid] = replace(s, updated_at=datetime.now(tz=timezone.utc) - timedelta(seconds=9999))
    r = execution.run_once()
    assert r.stale_pending == 0  # disabled => never warns