prediction_engine/python_bridge/tests/test_mmap_reader.py

"""Test mmap reader can parse ticks written by Rust."""
import struct
import tempfile
from pathlib import Path

from prediction_engine.python_bridge.mmap_reader import (
    HEADER_FMT, HEADER_SIZE, MAGIC, TICK_FMT, TICK_SIZE, CompactTick, MmapReader,
)


def _write_test_mmap(path: str, ticks: list):
    capacity = 100
    file_size = HEADER_SIZE + capacity * TICK_SIZE
    with open(path, "wb") as f:
        f.write(struct.pack(HEADER_FMT, MAGIC, capacity, len(ticks), 0))
        for venue, side, price, size, ts, mhash in ticks:
            f.write(struct.pack(TICK_FMT, venue, side, price, size, ts, mhash))
        remaining = file_size - HEADER_SIZE - len(ticks) * TICK_SIZE
        f.write(b"\x00" * remaining)


def test_read_header():
    with tempfile.NamedTemporaryFile(suffix=".mmap", delete=False) as f:
        path = f.name
    _write_test_mmap(path, [])
    reader = MmapReader(path)
    reader.open()
    ticks = reader.read_new_ticks()
    assert ticks == []
    reader.close()
    Path(path).unlink()


def test_read_ticks():
    with tempfile.NamedTemporaryFile(suffix=".mmap", delete=False) as f:
        path = f.name
    test_ticks = [
        (0, 0, 0.62, 100.0, 1709726400_000_000, 12345),
        (1, 1, 0.35, 200.0, 1709726401_000_000, 67890),
    ]
    _write_test_mmap(path, test_ticks)
    reader = MmapReader(path)
    reader.open()
    ticks = reader.read_new_ticks()
    assert len(ticks) == 2
    assert ticks[0].venue == "polymarket"
    assert ticks[0].price == 0.62
    assert ticks[1].venue == "kalshi"
    assert ticks[1].side == "no"
    reader.close()
    Path(path).unlink()