piker/tests/test_ems_state.py

# piker: trading gear for hackers
# Copyright (C) Tyler Goodlet (in stewardship for pikers)

# This program is free software: you can redistribute it and/or
# modify it under the terms of the GNU Affero General Public
# License as published by the Free Software Foundation, either
# version 3 of the License, or (at your option) any later version.

# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU Affero General Public License for more details.

# You should have received a copy of the GNU Affero General Public
# License along with this program. If not, see
# <https://www.gnu.org/licenses/>.

'''
Deterministic order-event sequencing without feeds or credentials.

'''

from decimal import Decimal

import pytest

from piker.clearing._codec import ProtocolError
from piker.clearing._schema import (
    BrokerdCancel,
    BrokerdError,
    BrokerdFill,
    BrokerdOrderAck,
    BrokerdStatus,
    Order,
    Status,
)
from piker.clearing._state import Effects, OrderState


def order(
    oid: str = 'one',
    account: str = 'paper',
) -> Order:
    '''
    Build live client intent with controllable dialog/account
    identity.

    '''
    return Order(
        oid=oid,
        account=account,
        symbol='btc.usdt.binance',
        price=Decimal('10'),
        size=1.0,
        action='buy',
        exec_mode='live',
    )


def ack(
    oid: str = 'one',
    account: str = 'paper',
    reqid: int|str = 1,
) -> BrokerdOrderAck:
    '''
    Build the acknowledgement establishing an account/provider ID.

    '''
    return BrokerdOrderAck(oid=oid, account=account, reqid=reqid)


def fill(
    fill_id: str = 'fill-1',
    account: str = 'paper',
    reqid: int|str = 1,
) -> BrokerdFill:
    '''
    Build an execution whose delivery may precede its
    acknowledgement.

    '''
    return BrokerdFill(
        time_ns=1,
        broker_time=1.0,
        account=account,
        broker='binance',
        reqid=reqid,
        fill_id=fill_id,
        size=0.5,
        price=10.0,
        action='buy',
    )


def closed() -> BrokerdStatus:
    '''
    Build a terminal provider status for the default test order.

    '''
    return BrokerdStatus(
        time_ns=1,
        reqid=1,
        account='paper',
        broker='binance',
        status='closed',
    )


def test_cancel_before_ack_once() -> None:
    '''
    A cancel can reach EMS while broker submission is suspended.
    Register intent, cancel twice, then repeat the ack. Only the
    first ack may produce a cancel, carrying the original account and
    oid. The earlier published snapshot must remain unchanged.

    '''
    state: OrderState = OrderState()
    initial: Status = state.register(order())
    assert not state.cancel('one').broker
    assert not state.cancel('one').broker
    effects: Effects = state.accept(ack())
    assert len(effects.broker) == 1
    cancel: BrokerdCancel = effects.broker[0]
    assert (cancel.oid, cancel.reqid, cancel.account) == (
        'one',
        1,
        'paper',
    )
    assert not state.accept(ack()).broker
    assert initial.cancel_called is False


@pytest.mark.parametrize('terminal_first', [True, False])
def test_early_duplicate_and_late_fills(
    terminal_first: bool,
) -> None:
    '''
    Native callbacks may precede an ack and a final fill may follow
    closed. Buffer both event orders, correlate with the ack, then
    repeat the fill and inject a distinct late execution. Each fill
    must be emitted once; neither may reopen the completed dialog.

    '''
    state: OrderState = OrderState()
    state.register(order())
    events: list[BrokerdStatus|BrokerdFill] = (
        [closed(), fill()] if terminal_first else [fill(), closed()]
    )
    event: BrokerdStatus|BrokerdFill
    for event in events:
        assert not state.accept(event).clients
    effects: Effects = state.accept(ack())
    responses: list[str] = []
    msg: Status
    for msg in effects.clients:
        responses.append(msg.resp)
    assert responses == (
        ['closed', 'fill'] if terminal_first else ['fill', 'closed']
    )
    assert not state.accept(fill()).clients
    assert state.accept(fill('fill-2')).clients[0].resp == 'fill'
    assert not state.active
    assert state.terminal['one'].status.resp == 'closed'


def test_closed_before_ack_cancels_no_finished_order() -> None:
    '''
    An order can finish while its ack is delayed and cancel is
    queued. Replay the pre-ack terminal status before generating
    effects so the delayed acknowledgement cannot cancel an already
    closed order.

    '''
    state: OrderState = OrderState()
    state.register(order())
    state.cancel('one')
    state.accept(closed())
    assert not state.accept(ack()).broker


def test_accounts_partition_request_ids() -> None:
    '''
    Separate broker sessions may both allocate request ID one.
    Register and acknowledge two accounts, then route their fills.
    Correlation must retain account ownership even when IDs collide.

    '''
    state: OrderState = OrderState()
    oid: str
    account: str
    for oid, account in [('one', 'paper'), ('two', 'other')]:
        state.register(order(oid, account))
        state.accept(ack(oid, account))
        effect: Effects = state.accept(fill(account=account))
        assert effect.clients[0].oid == oid


def test_pre_ack_capacity_is_bounded() -> None:
    '''
    A backend that never acknowledges must not grow EMS memory
    forever. Exceed a deliberately tiny buffer and assert a protocol
    failure preserves the already buffered execution instead of
    dropping it.

    '''
    state: OrderState = OrderState(max_early=1)
    state.accept(fill())
    with pytest.raises(ProtocolError, match='capacity'):
        state.accept(fill('fill-2'))
    assert state.early[('paper', 1)] == [fill()]


def test_tombstone_eviction_removes_correlation() -> None:
    '''
    Terminal dialogs have finite retention. Complete two orders with
    a one-entry limit; evicting the first must also remove every old
    request ID so a later event cannot target a missing dialog.

    '''
    state: OrderState = OrderState(max_terminal=1)
    oid: str
    reqid: int
    for oid, reqid in [('one', 1), ('two', 2)]:
        state.register(order(oid))
        state.accept(ack(oid, reqid=reqid))
        state.accept(closed().copy(update={'reqid': reqid}))
    assert tuple(state.terminal) == ('two',)
    assert state.reqids == {('paper', 2): 'two'}


def test_edit_before_ack_cannot_duplicate_submission() -> None:
    '''
    An edit with no known backend request ID could place a second
    order. Hold the ack and try an edit; rejection must leave the
    original request intact for subsequent acknowledgement and
    cancellation.

    '''
    state: OrderState = OrderState()
    initial: Status = state.register(order())
    with pytest.raises(
        ProtocolError,
        match='before acknowledgement',
    ):
        state.register(order().copy(update={'price': Decimal('11')}))
    assert state.active['one'] is initial


def test_edit_ack_keeps_late_fills_on_original_dialog() -> None:
    '''
    An edit may replace the provider request ID while executions for
    the old ID remain in flight. A delayed duplicate ack must not
    revert the cancel target, and either ID must still route fills to
    the same EMS order. Cancel once per accepted request identity.

    '''
    state: OrderState = OrderState()
    state.register(order())
    state.accept(ack())
    state.register(order().copy(update={'price': Decimal('11')}))
    state.accept(ack(reqid=2))
    state.accept(ack(reqid=1))
    assert state.active['one'].reqid == 2
    assert state.cancel('one').broker[0].reqid == 2
    assert not state.cancel('one').broker
    assert state.accept(fill(reqid=1)).clients[0].oid == 'one'
    effects: Effects = state.accept(fill('fill-2', reqid=2))
    assert effects.clients[0].oid == 'one'
    assert state.active['one'].reqid == 2
    state.accept(closed())
    assert 'one' in state.active


def test_startup_collision_does_not_register_second_order() -> None:
    '''
    Two existing orders cannot claim one account/request ID pair.
    Correlate a normal order, then inject a startup order with the
    same provider ID. Rejection leaves ``OrderState.active`` and
    ``OrderState.dialogs`` free of the rejected second order.

    '''
    from piker.clearing._schema import BrokerdOpenOrder

    state: OrderState = OrderState()
    state.register(order())
    state.accept(ack())
    with pytest.raises(ProtocolError, match='ownership'):
        state.accept(
            BrokerdOpenOrder(
                reqid=1,
                order=order('two'),
                broker='binance',
                time_ns=1,
            )
        )
    assert tuple(state.dialogs) == ('one',)
    assert tuple(state.active) == ('one',)


def test_submission_and_edit_rejections() -> None:
    '''
    A new submission error has no future ack to unblock cancellation.
    Its typed request identifies a terminal rejection; an edit error
    must retain the original live control ID so it can still cancel.
    Exercise both with identical reasons to prove request semantics,
    instead of parsing provider error strings, choose the transition.

    '''
    from piker.clearing._state import broker_request

    state: OrderState = OrderState()
    state.register(order())
    error: BrokerdError = BrokerdError(
        account='paper',
        broker='binance',
        reason='rejected',
        request=broker_request(order()),
    )
    assert state.accept(error).clients[0].resp == 'error'
    assert not state.active
    assert 'one' in state.terminal

    state = OrderState()
    state.register(order())
    state.accept(ack())
    error = error.copy(
        update={
            'request': broker_request(order(), 1),
        }
    )
    state.accept(error)
    assert state.cancel('one').broker[0].reqid == 1