diff --git a/kalshi/config.py b/kalshi/config.py index a2f46efa..4c0853e8 100644 --- a/kalshi/config.py +++ b/kalshi/config.py @@ -3,7 +3,9 @@ from __future__ import annotations import logging +from collections.abc import Callable from dataclasses import dataclass, field +from typing import Any from urllib.parse import urlparse import httpx @@ -44,6 +46,14 @@ class KalshiConfig: Requires the ``h2`` package (install ``httpx[http2]`` or ``h2``). limits: Custom ``httpx.Limits`` for connection pool tuning. ``None`` uses httpx defaults. + ws_ping_interval: Interval (s) between WS keepalive pings. Default 20. + ws_close_timeout: Time (s) to wait for graceful WS close handshake. + Default 5. + ws_json_loads: Optional callable used to parse WS frames. ``None`` + falls back to :func:`json.loads`. Set to e.g. ``orjson.loads`` + for ~2-3x faster recv-loop parsing on high-volume streams. + ws_json_dumps: Optional callable used to serialize outbound WS + commands. ``None`` falls back to :func:`json.dumps`. """ base_url: str = PRODUCTION_BASE_URL # trailing slash is stripped automatically @@ -57,6 +67,10 @@ class KalshiConfig: ws_max_retries: int = DEFAULT_WS_MAX_RETRIES http2: bool = False limits: httpx.Limits | None = None + ws_ping_interval: float = 20.0 + ws_close_timeout: float = 5.0 + ws_json_loads: Callable[[bytes | str], Any] | None = None + ws_json_dumps: Callable[[Any], bytes | str] | None = None def __post_init__(self) -> None: # Strip trailing slash to prevent double-slash in auth signing paths diff --git a/kalshi/ws/channels.py b/kalshi/ws/channels.py index 1072e9ad..85e9856d 100644 --- a/kalshi/ws/channels.py +++ b/kalshi/ws/channels.py @@ -6,6 +6,7 @@ import collections import json import logging +from collections.abc import Callable from typing import Any from websockets.exceptions import ConnectionClosed @@ -116,8 +117,14 @@ def __init__( connection: ConnectionManager, *, stash_maxlen: int = 1000, + json_loads: Callable[[bytes | str], Any] | None = None, ) -> None: self._connection = connection + # #209: pluggable JSON loader for subscribe-ack frames. ``None`` + # falls back to stdlib ``json.loads``. + self._json_loads: Callable[[bytes | str], Any] = ( + json_loads if json_loads is not None else json.loads + ) self._subscriptions: dict[int, Subscription] = {} # client_id -> Subscription self._sid_to_client: dict[int, int] = {} # server_sid -> client_id self._next_client_id = 1 @@ -195,7 +202,7 @@ async def _wait_for_response( raise KalshiConnectionError( f"Connection closed while awaiting response to command {msg_id}" ) from e - data: dict[str, Any] = json.loads(raw) + data: dict[str, Any] = self._json_loads(raw) if data.get("id") == msg_id: return data # Non-matching frame. During resubscribe, stash it by sid for diff --git a/kalshi/ws/client.py b/kalshi/ws/client.py index 6f2b8785..44087eeb 100644 --- a/kalshi/ws/client.py +++ b/kalshi/ws/client.py @@ -90,6 +90,11 @@ def __init__( self._running = False self._subscribe_lock = asyncio.Lock() self._pending_callbacks: list[tuple[str, Callable[..., Awaitable[None]]]] = [] + # #209: pluggable JSON loader (None -> stdlib json.loads). Resolved + # once here so the recv hot path doesn't keep dereferencing config. + self._json_loads: Callable[[bytes | str], Any] = ( + config.ws_json_loads if config.ws_json_loads is not None else json.loads + ) def connect(self) -> _WebSocketSession: """Return an async context manager for the WebSocket session.""" @@ -105,7 +110,7 @@ async def _start(self) -> None: ) await self._connection.connect() - self._sub_mgr = SubscriptionManager(self._connection) + self._sub_mgr = SubscriptionManager(self._connection, json_loads=self._json_loads) self._seq_tracker = SequenceTracker(on_gap=self._handle_seq_gap) self._orderbook_mgr = OrderbookManager() self._dispatcher = MessageDispatcher( @@ -138,10 +143,9 @@ async def _broadcast_sentinels(self) -> None: """Put a shutdown sentinel on every active subscription queue. Iterator consumers see the sentinel and exit their ``async for`` - loops. ``MessageQueue.put_sentinel`` is idempotent — calling - twice puts two sentinels in the queue, the iterator reads the - first and raises ``StopAsyncIteration`` without ever touching - the second. + loops. ``MessageQueue.put_sentinel`` is idempotent — once the + queue is closed, subsequent calls are no-ops (no extra sentinel + appended, no eviction of an in-flight item). Used by ``_stop()`` (post-cancel cleanup), the cooperative shutdown branch of ``run_forever(stop_event=...)`` (#177), and @@ -283,7 +287,7 @@ async def _process_frame(self, raw: str) -> None: desynced-and-still-consumed. """ assert self._dispatcher is not None - data = json.loads(raw) + data = self._json_loads(raw) sid = data.get("sid") seq = data.get("seq") msg_type = data.get("type", "") @@ -317,13 +321,18 @@ async def _process_frame(self, raw: str) -> None: pre_validated: BaseModel | None = None if msg_type == "orderbook_snapshot" and self._orderbook_mgr: snapshot = OrderbookSnapshotMessage.model_validate(data) + # #199: mutate in place — skip the O(n log n) sort + 2N + # OrderbookLevel allocation that the public ``apply_snapshot`` + # wrapper performs. Consumers materialize via ``get(ticker)`` + # on demand. # Record per-sid ticker ownership so all-markets subscriptions # can be torn down on gap recovery / unsubscribe (#189, #206). - self._orderbook_mgr.apply_snapshot(snapshot, sid=snapshot.sid) + self._orderbook_mgr._apply_snapshot_inplace(snapshot, sid=snapshot.sid) pre_validated = snapshot elif msg_type == "orderbook_delta" and self._orderbook_mgr: delta = OrderbookDeltaMessage.model_validate(data) - self._orderbook_mgr.apply_delta(delta) + # #199: in-place mutation only — see snapshot branch above. + self._orderbook_mgr._apply_delta_inplace(delta) pre_validated = delta try: diff --git a/kalshi/ws/connection.py b/kalshi/ws/connection.py index f7316d59..b5c0779c 100644 --- a/kalshi/ws/connection.py +++ b/kalshi/ws/connection.py @@ -64,6 +64,11 @@ def __init__( self._on_state_change = on_state_change self._ws: ClientConnection | None = None self._state = ConnectionState.DISCONNECTED + # #209: pluggable JSON dumper for outbound frames. Resolved once at + # init so the send hot path doesn't keep dereferencing config. + self._json_dumps: Callable[[Any], bytes | str] = ( + config.ws_json_dumps if config.ws_json_dumps is not None else json.dumps + ) @property def state(self) -> ConnectionState: @@ -107,6 +112,22 @@ async def _build_auth_headers(self) -> dict[str, str]: ws_path = urlparse(self._config.ws_base_url).path return await self._auth.sign_request_async("GET", ws_path) + async def _open_socket(self) -> ClientConnection: + """Open the websocket. Single source of truth for connect args. + + Both ``connect`` and ``reconnect`` route through this so the + ping/close timeouts and auth-header build stay in lockstep + (#208). + """ + headers = await self._build_auth_headers() + return await connect( + self._config.ws_base_url, + additional_headers=headers, + ping_interval=self._config.ws_ping_interval, + ping_timeout=self._heartbeat_timeout, + close_timeout=self._config.ws_close_timeout, + ) + async def connect(self) -> None: """Establish a WebSocket connection with RSA-PSS auth headers. @@ -117,14 +138,7 @@ async def connect(self) -> None: """ await self._set_state(ConnectionState.CONNECTING) try: - headers = await self._build_auth_headers() - self._ws = await connect( - self._config.ws_base_url, - additional_headers=headers, - ping_interval=20, - ping_timeout=self._heartbeat_timeout, - close_timeout=5.0, - ) + self._ws = await self._open_socket() await self._set_state(ConnectionState.CONNECTED) except Exception as e: await self._set_state(ConnectionState.CLOSED) @@ -139,11 +153,9 @@ async def connect(self) -> None: ) from e async def reconnect(self) -> None: - """Reconnect with exponential backoff and jitter. + """Reconnect with AWS Full Jitter backoff (matches REST transport). - Uses the same backoff pattern as the REST transport: - delay = retry_base_delay * (2 ** attempt) + random.uniform(0, 0.5) - delay = min(delay, retry_max_delay) + delay = random.uniform(0, min(retry_max_delay, retry_base_delay * 2 ** attempt)) Transitions: -> RECONNECTING -> CONNECTING -> CONNECTED or -> CLOSED (if max retries exceeded) @@ -159,10 +171,18 @@ async def reconnect(self) -> None: await self._set_state(ConnectionState.RECONNECTING) for attempt in range(self._config.ws_max_retries): - delay = self._config.retry_base_delay * ( - 2**attempt - ) + random.uniform(0, 0.5) - delay = min(delay, self._config.retry_max_delay) + # #221 P2.1: match REST transport's AWS Full Jitter formula. + # The previous ``base * 2**attempt + uniform(0, 0.5)`` reduced + # to a flat 0.5s jitter window once the exponential pinned the + # delay at ``retry_max_delay`` — thundering-herd risk on a + # global Kalshi WS outage. + delay = random.uniform( + 0, + min( + self._config.retry_max_delay, + self._config.retry_base_delay * (2**attempt), + ), + ) logger.warning( "Reconnecting in %.1fs (attempt %d/%d)", delay, @@ -172,14 +192,7 @@ async def reconnect(self) -> None: await asyncio.sleep(delay) try: await self._set_state(ConnectionState.CONNECTING) - headers = await self._build_auth_headers() - self._ws = await connect( - self._config.ws_base_url, - additional_headers=headers, - ping_interval=20, - ping_timeout=self._heartbeat_timeout, - close_timeout=5.0, - ) + self._ws = await self._open_socket() await self._set_state(ConnectionState.CONNECTED) return except Exception: @@ -212,7 +225,7 @@ async def send(self, msg: dict[str, Any]) -> None: """ if self._ws is None: raise KalshiConnectionError("Not connected") - await self._ws.send(json.dumps(msg)) + await self._ws.send(self._json_dumps(msg)) async def recv(self) -> str: """Receive a raw string message from the WebSocket. diff --git a/kalshi/ws/models/orderbook_delta.py b/kalshi/ws/models/orderbook_delta.py index 4bee6934..2eb1217b 100644 --- a/kalshi/ws/models/orderbook_delta.py +++ b/kalshi/ws/models/orderbook_delta.py @@ -1,6 +1,8 @@ """Orderbook delta and snapshot message models.""" from __future__ import annotations +from typing import Literal + from pydantic import AliasChoices, BaseModel, Field from kalshi.types import DollarDecimal, FixedPointCount @@ -45,7 +47,7 @@ class OrderbookDeltaPayload(BaseModel): delta: FixedPointCount = Field( validation_alias=AliasChoices("delta_fp", "delta"), ) - side: str + side: Literal["yes", "no"] client_order_id: str | None = None subaccount: int | None = None ts: str | None = None diff --git a/kalshi/ws/orderbook.py b/kalshi/ws/orderbook.py index 95056b26..3a5e3803 100644 --- a/kalshi/ws/orderbook.py +++ b/kalshi/ws/orderbook.py @@ -90,10 +90,14 @@ def __init__(self) -> None: # number of tickers owned by the sid, not O(n) over every book. self._sid_tickers: dict[int, set[str]] = {} - def apply_snapshot( + def _apply_snapshot_inplace( self, msg: OrderbookSnapshotMessage, sid: int | None = None - ) -> Orderbook: - """Initialize (or reset) a book from a full snapshot. + ) -> None: + """In-place state mutation for a snapshot. No materialization. + + Recv-loop hot path entry point (#199). The public + :meth:`apply_snapshot` wraps this with a ``to_orderbook()`` call + for direct callers that still want a fully built ``Orderbook``. ``sid`` records which subscription produced this snapshot so the ticker can be torn down later via :meth:`remove_by_sid` — required @@ -131,23 +135,19 @@ def apply_snapshot( len(no_levels), sid_val, ) - return state.to_orderbook() - - def apply_delta(self, msg: OrderbookDeltaMessage) -> Orderbook | None: - """Apply an incremental delta to an existing book. - O(1) on the update itself (price-indexed dict lookup); the returned - Orderbook is materialized at O(n log n) only because consumers expect - a sorted level list. + def _apply_delta_inplace(self, msg: OrderbookDeltaMessage) -> bool: + """In-place state mutation for a delta. No materialization. - Returns the updated Orderbook, or None if no book exists for this ticker - (delta arrived before snapshot -- should not happen in normal flow). + Returns True if the mutation was applied, False if no book exists + for the ticker yet (delta arrived before snapshot — caller may + choose to skip materializing in that case). """ ticker = msg.msg.market_ticker state = self._books.get(ticker) if state is None: logger.warning("Delta for unknown ticker %s (no snapshot yet)", ticker) - return None + return False price = msg.msg.price # Decimal via DollarDecimal delta = msg.msg.delta # Decimal via FixedPointCount @@ -165,8 +165,35 @@ def apply_delta(self, msg: OrderbookDeltaMessage) -> Orderbook | None: elif delta > 0: levels[price] = delta + return True + + def apply_snapshot( + self, msg: OrderbookSnapshotMessage, sid: int | None = None + ) -> Orderbook: + """Initialize (or reset) a book from a full snapshot. + + Public wrapper that mutates in place then materializes the + resulting :class:`Orderbook`. The recv loop bypasses this and + calls :meth:`_apply_snapshot_inplace` directly to avoid the + unused O(n log n) materialization (#199). + """ + self._apply_snapshot_inplace(msg, sid=sid) + ticker = msg.msg.market_ticker + state = self._books[ticker] return state.to_orderbook() + def apply_delta(self, msg: OrderbookDeltaMessage) -> Orderbook | None: + """Apply an incremental delta to an existing book. + + Public wrapper. Returns the materialized Orderbook on success or + None if no book exists for the ticker. The recv loop bypasses this + and calls :meth:`_apply_delta_inplace` directly (#199). + """ + if not self._apply_delta_inplace(msg): + return None + state = self._books.get(msg.msg.market_ticker) + return state.to_orderbook() if state else None + def get(self, ticker: str) -> Orderbook | None: """Get current book state (non-blocking). diff --git a/tests/test_config.py b/tests/test_config.py index f271beb4..b52a87c1 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -228,3 +228,36 @@ def test_total_timeout_field_exists_and_defaults_None(self) -> None: # noqa: N8 def test_total_timeout_can_be_set(self) -> None: config = KalshiConfig(total_timeout=30.0) assert config.total_timeout == 30.0 + + +class TestWsExtraConfigFields: + """#208 + #209: ws_ping_interval / ws_close_timeout / ws_json_loads / + ws_json_dumps fields exist on KalshiConfig with sane defaults. + """ + + def test_ws_ping_interval_default_20(self) -> None: + assert KalshiConfig().ws_ping_interval == 20.0 + + def test_ws_close_timeout_default_5(self) -> None: + assert KalshiConfig().ws_close_timeout == 5.0 + + def test_ws_json_loads_default_None(self) -> None: # noqa: N802 + assert KalshiConfig().ws_json_loads is None + + def test_ws_json_dumps_default_None(self) -> None: # noqa: N802 + assert KalshiConfig().ws_json_dumps is None + + def test_ws_ping_interval_settable(self) -> None: + cfg = KalshiConfig(ws_ping_interval=45.0) + assert cfg.ws_ping_interval == 45.0 + + def test_ws_close_timeout_settable(self) -> None: + cfg = KalshiConfig(ws_close_timeout=1.5) + assert cfg.ws_close_timeout == 1.5 + + def test_ws_json_loads_callable_stored(self) -> None: + import json + + cfg = KalshiConfig(ws_json_loads=json.loads, ws_json_dumps=json.dumps) + assert cfg.ws_json_loads is json.loads + assert cfg.ws_json_dumps is json.dumps diff --git a/tests/ws/test_backpressure.py b/tests/ws/test_backpressure.py index 58414abc..362c04b7 100644 --- a/tests/ws/test_backpressure.py +++ b/tests/ws/test_backpressure.py @@ -139,3 +139,26 @@ async def test_deque_maxlen_caps_memory_even_if_counter_drifts(self) -> None: await q.put(i) # Hard ceiling: maxsize + 1 (the +1 covers the put_sentinel append). assert len(q._buffer) <= q._maxsize + 1 + + +@pytest.mark.asyncio +class TestPutSentinelIdempotence: + """#221 P2.5: put_sentinel is a no-op when the queue is already closed. + + Previously the second call appended a second sentinel to the deque, and + when the buffer was at ``maxsize`` real items + 1 sentinel the deque's + ``maxlen=maxsize+1`` would evict the leftmost real item to make room. + """ + + async def test_put_sentinel_called_twice_is_idempotent(self) -> None: + q: MessageQueue[int] = MessageQueue(maxsize=3) + await q.put(1) + await q.put(2) + await q.put(3) + await q.put_sentinel() + before_len = len(q._buffer) + await q.put_sentinel() # MUST be a no-op + assert len(q._buffer) == before_len + # And the leftmost real item is still there. + items = [item async for item in q] + assert items == [1, 2, 3] diff --git a/tests/ws/test_client.py b/tests/ws/test_client.py index 9bc2b38c..081955d3 100644 --- a/tests/ws/test_client.py +++ b/tests/ws/test_client.py @@ -494,3 +494,42 @@ async def on_err(err: object) -> None: # Deterministic wait: the callback signals us; we don't sleep blindly. await asyncio.wait_for(got_one.wait(), timeout=2.0) assert len(errors) == 1 + + +# --------------------------------------------------------------------------- +# #209 — pluggable JSON loader +# --------------------------------------------------------------------------- + + +class TestPluggableJsonLoads: + async def test_custom_json_loads_called_for_recv_frames( + self, fake_ws: Any, test_auth: Any + ) -> None: + """#209: ws_json_loads is used for incoming WS frames in the recv loop.""" + import json as _json + + calls: list[Any] = [] + + def my_loads(raw: bytes | str) -> Any: + calls.append(raw) + return _json.loads(raw) + + config = KalshiConfig( + ws_base_url=fake_ws.url, timeout=5.0, ws_json_loads=my_loads + ) + ws = KalshiWebSocket(auth=test_auth, config=config) + async with ws.connect() as session: + stream = await session.subscribe_ticker(tickers=["T1"]) + await fake_ws.send_to_all( + { + "type": "ticker", + "sid": 1, + "msg": ticker_payload_dict( + market_ticker="T1", market_id="x", yes_bid_dollars="55" + ), + } + ) + await asyncio.wait_for(stream.__anext__(), timeout=2.0) + # At minimum the ticker frame was parsed by our loader. The recv loop + # parses every frame off the socket — assert non-empty. + assert calls, "ws_json_loads was never invoked on recv frames" diff --git a/tests/ws/test_connection.py b/tests/ws/test_connection.py index 3bbff39b..6e807c6d 100644 --- a/tests/ws/test_connection.py +++ b/tests/ws/test_connection.py @@ -471,3 +471,132 @@ async def test_received_commands_recorded( assert len(fake_ws.received_commands) == 1 assert fake_ws.received_commands[0]["cmd"] == "subscribe" await mgr.close() + + +# --------------------------------------------------------------------------- +# #208 — ws_ping_interval / ws_close_timeout pluming, #209 ws_json_dumps, +# #221 P2.1 Full Jitter +# --------------------------------------------------------------------------- + + +class TestPingCloseTimeoutPlumbing: + async def test_ws_ping_interval_from_config_passed_through( + self, fake_ws: FakeKalshiWS, test_auth: object, monkeypatch: pytest.MonkeyPatch + ) -> None: + """#208: ws_ping_interval forwarded into websockets.connect().""" + captured: dict[str, object] = {} + from kalshi.ws import connection as conn_mod + + real_connect = conn_mod.connect + + async def spy(*args: object, **kwargs: object) -> object: + captured.update(kwargs) + return await real_connect(*args, **kwargs) # type: ignore[misc] + + monkeypatch.setattr(conn_mod, "connect", spy) + + config = KalshiConfig( + ws_base_url=fake_ws.url, timeout=5.0, ws_ping_interval=42.5 + ) + mgr = ConnectionManager(auth=test_auth, config=config) # type: ignore[arg-type] + await mgr.connect() + try: + assert captured["ping_interval"] == 42.5 + finally: + await mgr.close() + + async def test_ws_close_timeout_from_config_passed_through( + self, fake_ws: FakeKalshiWS, test_auth: object, monkeypatch: pytest.MonkeyPatch + ) -> None: + """#208: ws_close_timeout forwarded into websockets.connect().""" + captured: dict[str, object] = {} + from kalshi.ws import connection as conn_mod + + real_connect = conn_mod.connect + + async def spy(*args: object, **kwargs: object) -> object: + captured.update(kwargs) + return await real_connect(*args, **kwargs) # type: ignore[misc] + + monkeypatch.setattr(conn_mod, "connect", spy) + + config = KalshiConfig( + ws_base_url=fake_ws.url, timeout=5.0, ws_close_timeout=2.25 + ) + mgr = ConnectionManager(auth=test_auth, config=config) # type: ignore[arg-type] + await mgr.connect() + try: + assert captured["close_timeout"] == 2.25 + finally: + await mgr.close() + + +class TestCustomJsonDumps: + async def test_custom_json_dumps_called_for_subscribe_commands( + self, fake_ws: FakeKalshiWS, test_auth: object + ) -> None: + """#209: ws_json_dumps used for outbound frame serialization.""" + calls: list[object] = [] + + def my_dumps(obj: object) -> str: + calls.append(obj) + return json.dumps(obj) + + config = KalshiConfig( + ws_base_url=fake_ws.url, timeout=5.0, ws_json_dumps=my_dumps + ) + mgr = ConnectionManager(auth=test_auth, config=config) # type: ignore[arg-type] + await mgr.connect() + try: + await mgr.send({"id": 1, "cmd": "subscribe", "params": {}}) + assert any( + isinstance(c, dict) and c.get("cmd") == "subscribe" for c in calls + ) + finally: + await mgr.close() + + +class TestReconnectFullJitter: + async def test_reconnect_uses_full_jitter( + self, test_auth: object, monkeypatch: pytest.MonkeyPatch + ) -> None: + """#221 P2.1: reconnect delay sampled via random.uniform(0, min(cap, base*2**n)) + — matching the REST transport's Full Jitter formula. Old formula was + base*2**n + uniform(0, 0.5) capped at retry_max_delay. + """ + from kalshi.ws import connection as conn_mod + + upper_bounds: list[float] = [] + + def fake_uniform(lo: float, hi: float) -> float: + upper_bounds.append(hi) + assert lo == 0 + return 0.0 # zero delay so the test doesn't sleep + + async def fake_sleep(_: float) -> None: + return None + + async def failing_connect(*_a: object, **_kw: object) -> object: + raise RuntimeError("simulated connect failure") + + monkeypatch.setattr(conn_mod.random, "uniform", fake_uniform) + monkeypatch.setattr(conn_mod.asyncio, "sleep", fake_sleep) + monkeypatch.setattr(conn_mod, "connect", failing_connect) + + config = KalshiConfig( + ws_base_url="ws://127.0.0.1:1", + timeout=1.0, + ws_max_retries=4, + retry_base_delay=0.5, + retry_max_delay=8.0, + ) + mgr = ConnectionManager(auth=test_auth, config=config) # type: ignore[arg-type] + with pytest.raises(KalshiConnectionError): + await mgr.reconnect() + + # Expected upper bounds per Full Jitter formula across 4 attempts: + # attempt 0: min(8.0, 0.5*1) = 0.5 + # attempt 1: min(8.0, 0.5*2) = 1.0 + # attempt 2: min(8.0, 0.5*4) = 2.0 + # attempt 3: min(8.0, 0.5*8) = 4.0 + assert upper_bounds == [0.5, 1.0, 2.0, 4.0] diff --git a/tests/ws/test_models.py b/tests/ws/test_models.py index d5acd0a4..6e564b3f 100644 --- a/tests/ws/test_models.py +++ b/tests/ws/test_models.py @@ -1129,3 +1129,63 @@ def test_communications_timestamps_parse_as_datetime(self) -> None: def test_rest_ws_field_type_symmetry(rest_field: object, ws_field: object) -> None: """REST and WS payloads MUST agree on the annotation for shared logical fields (#198).""" assert rest_field == ws_field + + +class TestOrderbookDeltaPayloadSideLiteral: + """#221 P2.2: OrderbookDeltaPayload.side typed as Literal['yes','no'].""" + + def test_orderbook_delta_payload_side_accepts_yes(self) -> None: + from kalshi.ws.models.orderbook_delta import OrderbookDeltaPayload + + payload = OrderbookDeltaPayload.model_validate( + { + "market_ticker": "T", + "market_id": "id", + "price": "0.50", + "delta": "10.00", + "side": "yes", + } + ) + assert payload.side == "yes" + + def test_orderbook_delta_payload_side_accepts_no(self) -> None: + from kalshi.ws.models.orderbook_delta import OrderbookDeltaPayload + + payload = OrderbookDeltaPayload.model_validate( + { + "market_ticker": "T", + "market_id": "id", + "price": "0.50", + "delta": "10.00", + "side": "no", + } + ) + assert payload.side == "no" + + def test_orderbook_delta_payload_side_rejects_invalid_string(self) -> None: + from kalshi.ws.models.orderbook_delta import OrderbookDeltaPayload + + with pytest.raises(ValidationError): + OrderbookDeltaPayload.model_validate( + { + "market_ticker": "T", + "market_id": "id", + "price": "0.50", + "delta": "10.00", + "side": "maybe", + } + ) + + def test_orderbook_delta_payload_side_rejects_trailing_whitespace(self) -> None: + from kalshi.ws.models.orderbook_delta import OrderbookDeltaPayload + + with pytest.raises(ValidationError): + OrderbookDeltaPayload.model_validate( + { + "market_ticker": "T", + "market_id": "id", + "price": "0.50", + "delta": "10.00", + "side": "yes ", + } + ) diff --git a/tests/ws/test_orderbook.py b/tests/ws/test_orderbook.py index 42e40ea7..dac70edc 100644 --- a/tests/ws/test_orderbook.py +++ b/tests/ws/test_orderbook.py @@ -268,3 +268,99 @@ def test_fractional_delta_matches_wire_format(self) -> None: # Decimal("0.00") + Decimal("1.00") == Decimal("1.00"), not Decimal("1") assert book.yes[0].quantity == Decimal("1.00") assert book.yes[0].price == Decimal("0.0100") + + +class TestOrderbookManagerInplace: + """#199: in-place variants used by the recv loop to avoid the + O(n log n) sort + 2N OrderbookLevel allocation that the public + ``apply_*`` wrappers perform. + """ + + def test_apply_snapshot_inplace_returns_none_and_mutates(self) -> None: + mgr = OrderbookManager() + snap = make_snapshot( + yes=[["0.50", "100.00"], ["0.55", "200.00"]], + no=[["0.45", "50.00"]], + ) + result = mgr._apply_snapshot_inplace(snap) + assert result is None + # State should be mutated; ``get`` materializes on demand. + book = mgr.get("T") + assert book is not None + assert len(book.yes) == 2 + assert len(book.no) == 1 + + def test_apply_snapshot_inplace_records_sid(self) -> None: + mgr = OrderbookManager() + snap = make_snapshot(yes=[["0.50", "100.00"]]) + mgr._apply_snapshot_inplace(snap, sid=snap.sid) + assert mgr.tickers_for_sid(snap.sid) == {"T"} + + def test_apply_delta_inplace_returns_true_on_known_ticker(self) -> None: + mgr = OrderbookManager() + mgr._apply_snapshot_inplace(make_snapshot(yes=[["0.50", "10.00"]])) + ok = mgr._apply_delta_inplace( + make_delta(price="0.50", delta="5.00", side="yes") + ) + assert ok is True + book = mgr.get("T") + assert book is not None + assert book.yes[0].quantity == Decimal("15.00") + + def test_apply_delta_inplace_returns_false_for_unknown_ticker(self) -> None: + mgr = OrderbookManager() + ok = mgr._apply_delta_inplace( + make_delta(ticker="OTHER", price="0.50", delta="5.00") + ) + assert ok is False + + def test_apply_delta_inplace_does_not_call_to_orderbook( + self, monkeypatch: object, + ) -> None: + """Hot-path invariant: the in-place variant MUST NOT materialize + an Orderbook. Spy on ``_BookState.to_orderbook`` and assert zero + calls during a snapshot+delta cycle through the in-place API. + """ + from kalshi.ws import orderbook as ob_mod + + calls: list[str] = [] + orig = ob_mod._BookState.to_orderbook + + def spy(self): # type: ignore[no-untyped-def] + calls.append(self.ticker) + return orig(self) + + monkeypatch.setattr(ob_mod._BookState, "to_orderbook", spy) # type: ignore[attr-defined] + + mgr = OrderbookManager() + mgr._apply_snapshot_inplace(make_snapshot(yes=[["0.50", "10.00"]])) + for _ in range(5): + mgr._apply_delta_inplace( + make_delta(price="0.50", delta="1.00", side="yes") + ) + assert calls == [] + + def test_public_apply_delta_still_returns_orderbook(self) -> None: + """Back-compat: direct callers of the public ``apply_delta`` still + get a fully materialized Orderbook. + """ + mgr = OrderbookManager() + mgr.apply_snapshot(make_snapshot(yes=[["0.50", "10.00"]])) + book = mgr.apply_delta(make_delta(price="0.50", delta="1.00", side="yes")) + assert book is not None + assert book.yes[0].quantity == Decimal("11.00") + + def test_public_apply_delta_returns_none_for_unknown_ticker(self) -> None: + mgr = OrderbookManager() + result = mgr.apply_delta(make_delta(ticker="UNK")) + assert result is None + + def test_public_apply_snapshot_returns_orderbook(self) -> None: + mgr = OrderbookManager() + book = mgr.apply_snapshot( + make_snapshot(yes=[["0.10", "5.00"]], no=[["0.90", "5.00"]]) + ) + assert book is not None + assert book.ticker == "T" + assert len(book.yes) == 1 + assert len(book.no) == 1