poly-maker/tests/test_state.py

"""Tests for StateStore, the user-event tracker, and the reconciler."""

from __future__ import annotations

from polymaker.domain import (
    Fill,
    OpenOrder,
    OrderState,
    Quote,
    Regime,
    Side,
    TargetQuotes,
    TradeState,
)
from polymaker.execution.reconciler import reconcile
from polymaker.state.store import StateStore
from polymaker.state.tracker import OrderEvent, TradeEvent, UserEventProcessor

# ── StateStore ──────────────────────────────────────────────────────────────


def test_apply_fill_updates_size_and_avg(tmp_path):
    s = StateStore(tmp_path / "s.db")
    s.apply_fill(Fill("tok", Side.BUY, 0.50, 100, "t1"))
    assert s.position("tok").size == 100
    assert s.position("tok").avg_price == 0.50
    # buy more at a higher price -> weighted avg
    s.apply_fill(Fill("tok", Side.BUY, 0.60, 100, "t2"))
    assert s.position("tok").size == 200
    assert abs(s.position("tok").avg_price - 0.55) < 1e-9
    # sell reduces size, avg unchanged
    s.apply_fill(Fill("tok", Side.SELL, 0.70, 50, "t3"))
    assert s.position("tok").size == 150
    assert abs(s.position("tok").avg_price - 0.55) < 1e-9
    s.close()


def test_sell_to_flat_resets_avg(tmp_path):
    s = StateStore(tmp_path / "s.db")
    s.apply_fill(Fill("tok", Side.BUY, 0.5, 100, "t1"))
    s.apply_fill(Fill("tok", Side.SELL, 0.6, 100, "t2"))
    assert s.position("tok").size == 0
    assert s.position("tok").avg_price == 0.0
    s.close()


def test_reconcile_positions_skips_inflight_and_recent(tmp_path):
    s = StateStore(tmp_path / "s.db")
    s.mark_inflight("tok")
    s.reconcile_positions({"tok": (999.0, 0.9)})  # ignored: in-flight
    assert s.position("tok").size == 0
    s.clear_inflight("tok")
    # still recent fill guard: simulate no recent fill by using a fresh token
    s.reconcile_positions({"other": (42.0, 0.3)})
    assert s.position("other").size == 42.0
    s.close()


def test_state_persists_across_restart(tmp_path):
    db = tmp_path / "s.db"
    s = StateStore(db)
    s.apply_fill(Fill("tok", Side.BUY, 0.5, 100, "t1"))
    s.close()
    s2 = StateStore(db)
    assert s2.position("tok").size == 100
    s2.close()


# ── UserEventProcessor ───────────────────────────────────────────────────────


def test_matched_then_confirmed(tmp_path):
    s = StateStore(tmp_path / "s.db")
    changed: list[str] = []
    p = UserEventProcessor(s, on_change=changed.append)
    p.on_trade(TradeEvent("tok", Side.BUY, 0.5, 100, "trade1", TradeState.MATCHED, 1.0), "cid")
    assert s.position("tok").size == 100
    assert s.inflight("tok") == 1
    p.on_trade(TradeEvent("tok", Side.BUY, 0.5, 100, "trade1", TradeState.CONFIRMED, 2.0), "cid")
    assert s.inflight("tok") == 0
    assert s.position("tok").size == 100  # settled
    assert changed == ["cid", "cid"]
    s.close()


def test_matched_is_idempotent(tmp_path):
    s = StateStore(tmp_path / "s.db")
    p = UserEventProcessor(s)
    ev = TradeEvent("tok", Side.BUY, 0.5, 100, "trade1", TradeState.MATCHED, 1.0)
    p.on_trade(ev, "cid")
    p.on_trade(ev, "cid")  # duplicate MATCHED for same trade id
    assert s.position("tok").size == 100  # not doubled
    s.close()


def test_failed_trade_reverses_fill(tmp_path):
    s = StateStore(tmp_path / "s.db")
    p = UserEventProcessor(s)
    p.on_trade(TradeEvent("tok", Side.BUY, 0.5, 100, "trade1", TradeState.MATCHED, 1.0), "cid")
    assert s.position("tok").size == 100
    p.on_trade(TradeEvent("tok", Side.BUY, 0.5, 100, "trade1", TradeState.FAILED, 2.0), "cid")
    assert s.position("tok").size == 0  # rolled back
    assert s.inflight("tok") == 0
    s.close()


def test_order_event_upsert_and_cancel(tmp_path):
    s = StateStore(tmp_path / "s.db")
    p = UserEventProcessor(s)
    p.on_order(OrderEvent("o1", "tok", Side.BUY, 0.49, 100), "cid")
    assert len(s.orders_for("tok")) == 1
    p.on_order(OrderEvent("o1", "tok", Side.BUY, 0.49, 0, is_cancel=True), "cid")
    assert len(s.orders_for("tok")) == 0
    s.close()


# ── reconciler ───────────────────────────────────────────────────────────────


def _live(order_id, token, side, price, size):
    return OpenOrder(order_id, token, side, price, size, OrderState.LIVE)


def test_reconcile_places_when_no_live():
    tq = TargetQuotes("cid", Regime.QUIET, (Quote("tok", Side.BUY, 0.49, 100),))
    plan = reconcile(tq, [], tick=0.01, reprice_ticks=2, resize_frac=0.15)
    assert len(plan.to_place) == 1
    assert plan.to_cancel == []


def test_reconcile_keeps_close_order():
    tq = TargetQuotes("cid", Regime.QUIET, (Quote("tok", Side.BUY, 0.49, 100),))
    live = [_live("o1", "tok", Side.BUY, 0.49, 102)]  # within tolerances
    plan = reconcile(tq, live, tick=0.01, reprice_ticks=2, resize_frac=0.15)
    assert plan.is_noop


def test_reconcile_reprices_when_far():
    tq = TargetQuotes("cid", Regime.QUIET, (Quote("tok", Side.BUY, 0.45, 100),))
    live = [_live("o1", "tok", Side.BUY, 0.49, 100)]  # 4 ticks away > 2
    plan = reconcile(tq, live, tick=0.01, reprice_ticks=2, resize_frac=0.15)
    assert plan.to_cancel == ["o1"]
    assert len(plan.to_place) == 1


def test_reconcile_resizes_when_size_drifts():
    tq = TargetQuotes("cid", Regime.QUIET, (Quote("tok", Side.BUY, 0.49, 100),))
    live = [_live("o1", "tok", Side.BUY, 0.49, 50)]  # 50% smaller > 15%
    plan = reconcile(tq, live, tick=0.01, reprice_ticks=2, resize_frac=0.15)
    assert plan.to_cancel == ["o1"]
    assert len(plan.to_place) == 1


def test_reconcile_cancels_all_when_target_empty():
    tq = TargetQuotes("cid", Regime.EVENT, ())
    live = [_live("o1", "tok", Side.BUY, 0.49, 100), _live("o2", "tok", Side.SELL, 0.55, 50)]
    plan = reconcile(tq, live, tick=0.01, reprice_ticks=2, resize_frac=0.15)
    assert set(plan.to_cancel) == {"o1", "o2"}
    assert plan.to_place == []


def test_reconcile_matches_layers_one_to_one():
    tq = TargetQuotes("cid", Regime.QUIET, (
        Quote("tok", Side.BUY, 0.49, 100),
        Quote("tok", Side.BUY, 0.47, 100),
    ))
    live = [_live("o1", "tok", Side.BUY, 0.49, 100)]  # only the top layer exists
    plan = reconcile(tq, live, tick=0.01, reprice_ticks=2, resize_frac=0.15)
    assert plan.to_cancel == []
    assert len(plan.to_place) == 1  # only the missing deeper layer
    assert plan.to_place[0].price == 0.47