trex-core/tests/application/test_venue_intent_flow.py

"""Venue propagation from the strategy wire contract to broker execution."""
from __future__ import annotations

from uuid import UUID

import pytest

from trex_core.adapters.persistence.in_memory_intent_store import InMemoryIntentStore
from trex_core.application.broker_port import (
    BalanceItem,
    CancelOutcome,
    CancelState,
    FillStatus,
    OrderAck,
    OrderRequest,
)
from trex_core.application.execution import ExecutionService
from trex_core.application.intents import BuyIntent, IntentService
from trex_core.application.messaging import RpcRequest
from trex_core.application.messaging_handlers import make_intent_rpc_handler
from trex_core.domain.slot import SlotService, SlotState
from trex_core.domain.slot.in_memory_repository import InMemorySlotRepository


class RecordingBroker:
    def __init__(self, *, actual_venue: str | None = None) -> None:
        self.orders: list[OrderRequest] = []
        self.cancel_venues: list[str] = []
        self.actual_venue = actual_venue

    def submit_order(self, req: OrderRequest) -> OrderAck:
        self.orders.append(req)
        return OrderAck(
            accepted=True,
            broker_order_id=f"order-{len(self.orders)}",
            venue=self.actual_venue or req.venue,
        )

    def poll_fill(self, broker_order_id: str, *, side: str, symbol: str) -> FillStatus:
        return FillStatus(state="accepted")

    def cancel_order(
        self,
        broker_order_id: str,
        *,
        side: str,
        symbol: str,
        qty: int,
        venue: str = "KRX",
    ) -> CancelOutcome:
        self.cancel_venues.append(venue)
        return CancelOutcome(CancelState.REQUESTED)

    def get_balance(self) -> list[BalanceItem]:
        return []


def _wire() -> tuple[
    InMemorySlotRepository,
    IntentService,
    ExecutionService,
    RecordingBroker,
]:
    repo = InMemorySlotRepository()
    service = SlotService(repo)
    intents = IntentService(
        service,
        repo,
        InMemoryIntentStore(),
        resolve_symbol=lambda portfolio_id: "005930",
    )
    repo.reconcile_universe("default", {"PF-1": "005930"})
    broker = RecordingBroker()
    return repo, intents, ExecutionService(service, repo, broker), broker


def test_nxt_buy_wire_value_reaches_broker_order() -> None:
    repo, intents, execution, broker = _wire()
    handler = make_intent_rpc_handler(intents)

    response = handler(
        RpcRequest(
            payload={
                "intent_type": "buy",
                "intent_id": "nxt-buy",
                "portfolio_id": "PF-1",
                "qty": 1,
                "venue": "NXT",
            },
            correlation_id="nxt-buy",
        )
    )

    assert response.payload["status"] == "admitted"
    pending = next(slot for slot in repo.list_slots("default") if slot.status is SlotState.PENDING_BUY)
    assert pending.venue == "NXT"
    execution.run_once()
    assert broker.orders[0].venue == "NXT"


@pytest.mark.parametrize("invalid_venue", ["AUTO", "UNKNOWN"])
def test_invalid_venue_is_rejected_before_slot_transition(invalid_venue: str) -> None:
    repo, intents, _, _ = _wire()
    handler = make_intent_rpc_handler(intents)

    response = handler(
        RpcRequest(
            payload={
                "intent_type": "buy",
                "intent_id": "bad-venue",
                "portfolio_id": "PF-1",
                "qty": 1,
                "venue": invalid_venue,
            },
            correlation_id="bad-venue",
        )
    )

    assert response.payload["status"] == "rejected"
    assert response.payload["reason"] == "invalid_venue"
    assert repo.list_slots("default")[0].status is SlotState.INACTIVE


def test_omitted_wire_venue_is_sor_then_actual_venue_is_persisted() -> None:
    repo = InMemorySlotRepository()
    service = SlotService(repo)
    intents = IntentService(
        service,
        repo,
        InMemoryIntentStore(),
        resolve_symbol=lambda portfolio_id: "005930",
    )
    repo.reconcile_universe("default", {"PF-1": "005930"})
    broker = RecordingBroker(actual_venue="NXT")
    execution = ExecutionService(service, repo, broker)
    handler = make_intent_rpc_handler(intents)

    response = handler(
        RpcRequest(
            payload={
                "intent_type": "buy",
                "intent_id": "auto-buy",
                "portfolio_id": "PF-1",
                "qty": 1,
            },
            correlation_id="auto-buy",
        )
    )

    slot_id = UUID(response.payload["slot_id"])
    pending = repo.fetch_slot(slot_id)
    assert pending is not None and pending.venue == "SOR"
    execution.run_once()
    dispatched = repo.fetch_slot(slot_id)
    assert broker.orders[0].venue == "SOR"
    assert dispatched is not None and dispatched.venue == "NXT"


def test_inactive_transition_resets_venue_to_krx() -> None:
    repo, intents, _, _ = _wire()
    handler = make_intent_rpc_handler(intents)
    handler(
        RpcRequest(
            payload={
                "intent_type": "buy",
                "intent_id": "nxt-buy-reset",
                "portfolio_id": "PF-1",
                "qty": 1,
                "venue": "NXT",
            },
            correlation_id="nxt-buy-reset",
        )
    )
    slot = repo.list_slots("default")[0]
    service = SlotService(repo)
    reset = service.transit_slot(
        slot.slot_id,
        SlotState.PENDING_BUY,
        SlotState.INACTIVE,
        reason="rollback",
        actor="test",
    )
    assert reset.venue == "KRX"


def test_holding_transition_resets_next_order_venue_to_sor() -> None:
    repo, intents, _, _ = _wire()
    decision = intents.submit_buy(
        BuyIntent(
            intent_id="nxt-buy-holding",
            portfolio_id="PF-1",
            qty=1,
            venue="NXT",
        )
    )
    slot = repo.fetch_slot(decision.slot_id)
    service = SlotService(repo)
    holding = service.transit_slot(
        slot.slot_id,
        SlotState.PENDING_BUY,
        SlotState.HOLDING,
        reason="fill",
        actor="test",
        qty=1,
        avg_price=1000,
    )
    assert holding.venue == "SOR"


def test_operator_cancel_uses_original_nxt_venue() -> None:
    repo, intents, execution, broker = _wire()
    handler = make_intent_rpc_handler(intents)
    response = handler(
        RpcRequest(
            payload={
                "intent_type": "buy",
                "intent_id": "nxt-buy-cancel",
                "portfolio_id": "PF-1",
                "qty": 1,
                "venue": "NXT",
            },
            correlation_id="nxt-buy-cancel",
        )
    )
    execution.run_once()

    assert execution.cancel_pending_slot(response.payload["slot_id"]) == "ok"
    assert broker.cancel_venues == ["NXT"]