trex-core/trex_core/adapters/broker/paper.py

"""Paper/stub broker — deterministic async double (REQ-011, dcy-003 OQ-2).

Follows the same submit -> poll contract as the real KIS adapter (no
synchronous shortcut), but lets tests program WHEN/whether fills happen:
immediate, after N polls, partial, or reject. This exercises the same
execution code path that production (KIS) takes.
"""
from __future__ import annotations

import threading
from decimal import Decimal
from typing import Optional

from trex_core.application.broker_port import (
    BalanceItem,
    FillStatus,
    OrderAck,
    OrderRequest,
    calculate_fee,
    calculate_tax,
)


class _Order:
    def __init__(self, oid: str, req: OrderRequest) -> None:
        self.oid = oid
        self.side = req.side
        self.symbol = req.symbol
        self.qty = req.qty
        self.price = float(req.limit_price) if req.limit_price is not None else 1000.0
        self.polls = 0
        self.filled = 0
        self.cancelled = False


class PaperBroker:
    """Deterministic async paper broker.

    - fill_after_polls: number of poll_fill calls before fills begin (0 => fills on first poll)
    - qty_per_poll: filled quantity per poll once filling (None => fill remaining at once)
    - reject_symbols: symbols whose orders are rejected on first poll
    """

    def __init__(
        self,
        *,
        fill_after_polls: int = 0,
        qty_per_poll: Optional[int] = None,
        reject_symbols: Optional[set[str]] = None,
    ) -> None:
        self._fill_after = fill_after_polls
        self._qty_per_poll = qty_per_poll
        self._reject = reject_symbols or set()
        self._orders: dict[str, _Order] = {}
        self._seq = 0
        self._lock = threading.Lock()
        # REQ-038: programmable current price for safety-sell scenario tests.
        self._prices: dict[str, Decimal] = {}

    def name(self) -> str:
        return "paper"

    def submit_order(self, req: OrderRequest) -> OrderAck:
        with self._lock:
            self._seq += 1
            oid = f"PAPER-{self._seq:06d}"
            self._orders[oid] = _Order(oid, req)
        return OrderAck(accepted=True, broker_order_id=oid)

    def poll_fill(self, broker_order_id: str, *, side: str, symbol: str) -> FillStatus:
        with self._lock:
            o = self._orders.get(broker_order_id)
            if o is None:
                return FillStatus(state="unknown")
            if o.cancelled:
                return FillStatus(state="cancelled")
            o.polls += 1
            if symbol in self._reject:
                return FillStatus(state="rejected", raw={"reason": "paper_reject"})
            if o.polls <= self._fill_after:
                return FillStatus(state="accepted", filled_qty=o.filled, avg_price=o.price)
            step = self._qty_per_poll if self._qty_per_poll is not None else (o.qty - o.filled)
            o.filled = min(o.qty, o.filled + step)
            if o.filled >= o.qty:
                # Deterministic execution_id (dcy-004 OQ-3):
                #   paper:{broker_order_id}:{filled_qty}
                execution_id = f"paper:{o.oid}:{o.filled}"
                amount = o.filled * o.price
                fee = calculate_fee(o.side, amount)
                tax = calculate_tax(o.side, amount)
                return FillStatus(
                    state="filled", filled_qty=o.filled, avg_price=o.price,
                    execution_id=execution_id, fee=fee, tax=tax,
                )
            return FillStatus(state="partially_filled", filled_qty=o.filled, avg_price=o.price)

    def cancel_order(self, broker_order_id: str, *, side: str, symbol: str, qty: int) -> bool:
        with self._lock:
            o = self._orders.get(broker_order_id)
            if o is None or o.filled >= o.qty:
                return False
            o.cancelled = True
            return True

    def get_balance(self) -> list[BalanceItem]:
        """Deterministic per-symbol fill aggregate (REQ-031 AC-3).

        Aggregates buys − sells across all filled (or partially filled)
        paper orders, drops zero-qty rows (Design Decision 7).
        """
        with self._lock:
            balances: dict[str, int] = {}
            for o in self._orders.values():
                if o.filled <= 0:
                    continue
                sign = 1 if o.side == "buy" else -1
                balances[o.symbol] = balances.get(o.symbol, 0) + sign * int(o.filled)
        return [BalanceItem(symbol=sym, hldg_qty=qty)
                for sym, qty in balances.items() if qty != 0]

    # --- REQ-038: programmable current price (safety-sell scenario double) ---
    def set_current_price(self, symbol: str, price: Decimal) -> None:
        """Program the current price returned by get_current_price(symbol)."""
        with self._lock:
            self._prices[symbol] = Decimal(str(price))

    def get_current_price(self, symbol: str) -> Optional[Decimal]:
        """Return the programmed current price for symbol, or None if unset."""
        with self._lock:
            return self._prices.get(symbol)