tests/test_orderbook.py

import pytest

from src.microstructure.orderbook import KalshiOrderBook
from src.microstructure.collector import TickCollector


# --- order-book reconstruction -------------------------------------------

def test_snapshot_top_of_book():
    book = KalshiOrderBook("TEST")
    book.apply_snapshot(yes=[[40, 100], [38, 50]], no=[[55, 80], [50, 20]])
    feats = book.features()
    assert feats["best_bid"] == 0.40          # highest yes bid
    assert feats["best_ask"] == 0.45          # 100 - highest no bid (55)
    assert feats["bid_size"] == 100
    assert feats["ask_size"] == 80
    assert feats["spread"] == pytest.approx(0.05)
    assert feats["mid"] == pytest.approx(0.425)


def test_delta_adds_and_removes_levels():
    book = KalshiOrderBook("TEST")
    book.apply_snapshot(yes=[[40, 100]], no=[[55, 80]])
    book.apply_delta(40, -100, "yes")         # wipe the only yes level
    assert book.features() is None            # one-sided -> no features
    book.apply_delta(42, 30, "yes")           # add a new level
    assert book.yes_bids == {42: 30}
    assert book.features()["best_bid"] == 0.42


def test_delta_partial_fill():
    book = KalshiOrderBook("TEST")
    book.apply_snapshot(yes=[[40, 100]], no=[[55, 80]])
    book.apply_delta(40, -60, "yes")
    assert book.yes_bids[40] == 40


def test_crossed_book_returns_no_features():
    book = KalshiOrderBook("TEST")
    # yes bid 60, no bid 65 -> ask 35 -> bid(60) > ask(35): crossed
    book.apply_snapshot(yes=[[60, 10]], no=[[65, 10]])
    assert book.features() is None


def test_imbalance_and_microprice():
    book = KalshiOrderBook("TEST")
    # heavy bid side: 300 resting bid vs 100 ask
    book.apply_snapshot(yes=[[40, 300]], no=[[55, 100]])  # ask = 45
    feats = book.features()
    assert feats["imbalance"] == 0.75                     # 300 / (300 + 100)
    # microprice = (40*100 + 45*300) / 400 = 43.75c, pulled toward the ask
    assert feats["microprice"] == pytest.approx(0.4375)


# --- collector message handling ------------------------------------------

def test_collector_applies_snapshot_then_delta():
    c = TickCollector(n_markets=1)
    snap = {"type": "orderbook_snapshot", "seq": 1,
            "msg": {"market_ticker": "T", "yes": [[40, 100]], "no": [[55, 80]]}}
    assert c._apply_message(snap) is False
    assert c.buffer[-1]["best_bid"] == 0.40

    delta = {"type": "orderbook_delta", "seq": 2,
             "msg": {"market_ticker": "T", "price": 42, "delta": 30, "side": "yes"}}
    assert c._apply_message(delta) is False
    assert c.buffer[-1]["best_bid"] == 0.42


def test_collector_detects_seq_gap():
    c = TickCollector(n_markets=1)
    c._apply_message({"type": "orderbook_snapshot", "seq": 1,
                      "msg": {"market_ticker": "T",
                              "yes": [[40, 100]], "no": [[55, 80]]}})
    gap = {"type": "orderbook_delta", "seq": 5,   # jumped 1 -> 5
           "msg": {"market_ticker": "T", "price": 42, "delta": 30, "side": "yes"}}
    assert c._apply_message(gap) is True


def test_collector_ignores_control_frames():
    c = TickCollector(n_markets=1)
    assert c._apply_message({"type": "subscribed", "msg": {"sid": 1}}) is False
    assert c.buffer == []