Proptest Invariant

// tests/unit/common/strategies.rs

use orderbook_rs::prelude::*;
use proptest::prelude::*;
use proptest::collection::vec;

/// Small, deterministic set of owner ids. Keeping the pool small makes STP / self-cross
/// exercise frequently on short streams.
const OWNER_POOL: [u64; 4] = [1, 2, 3, 4];

/// Operation generated by the strategy. Keep this enum local to the test crate so it
/// does not get re-exported by accident.
#[derive(Clone, Debug)]
pub enum Op {
    Submit {
        id: Id,
        owner: u64,
        side: Side,
        price: Price,
        qty: Quantity,
        tif: TimeInForce,
    },
    CancelById(Id),
    MassCancelBySide(Side),
}

pub fn op_stream(len_range: std::ops::Range<usize>) -> impl Strategy<Value = Vec<Op>> {
    vec(op(), len_range)
}

fn op() -> impl Strategy<Value = Op> {
    prop_oneof![
        7 => submit(),
        2 => cancel(),
        1 => mass_cancel(),
    ]
}

fn submit() -> impl Strategy<Value = Op> {
    (
        any::<u64>().prop_map(Id::from_u64),
        any::<usize>().prop_map(|i| OWNER_POOL[i % OWNER_POOL.len()]),
        any::<Side>(),
        // Tight price band forces crossings. A wide uniform range produces a flat empty
        // book that never exercises matching.
        (99u64..=101).prop_map(|t| Price::from_u64(t)),
        (1u64..=100).prop_map(|q| Quantity::from_u64(q)),
        tif_any(),
    )
        .prop_map(|(id, owner, side, price, qty, tif)| Op::Submit {
            id, owner, side, price, qty, tif,
        })
}

fn cancel() -> impl Strategy<Value = Op> {
    any::<u64>().prop_map(|v| Op::CancelById(Id::from_u64(v)))
}

fn mass_cancel() -> impl Strategy<Value = Op> {
    prop_oneof![
        Just(Op::MassCancelBySide(Side::Buy)),
        Just(Op::MassCancelBySide(Side::Sell)),
    ]
}

fn tif_any() -> impl Strategy<Value = TimeInForce> {
    prop_oneof![
        Just(TimeInForce::Gtc),
        Just(TimeInForce::Ioc),
        Just(TimeInForce::Fok),
        Just(TimeInForce::PostOnly),
    ]
}

// tests/unit/props_qty_conserved.rs

mod common;
use common::strategies::{op_stream, Op};

use orderbook_rs::prelude::*;
use proptest::prelude::*;

fn apply(book: &OrderBook<()>, op: &Op) {
    match op {
        Op::Submit { id, owner, side, price, qty, tif } => {
            let _ = book.submit_limit(*id, *owner, *side, *price, *qty, *tif);
        }
        Op::CancelById(id) => {
            let _ = book.cancel(*id);
        }
        Op::MassCancelBySide(side) => {
            let _ = book.mass_cancel_by_side(*side);
        }
    }
}

fn bid_qty_sum(book: &OrderBook<()>) -> u64 {
    book.iter_levels(Side::Buy)
        .map(|lvl| lvl.total_quantity().as_u64())
        .sum()
}

fn ask_qty_sum(book: &OrderBook<()>) -> u64 {
    book.iter_levels(Side::Sell)
        .map(|lvl| lvl.total_quantity().as_u64())
        .sum()
}

proptest! {
    #![proptest_config(ProptestConfig {
        cases: 256,
        max_shrink_iters: 50_000,
        ..ProptestConfig::default()
    })]

    #[test]
    fn qty_conserved_across_snapshot(stream in op_stream(1..200)) {
        let book = OrderBook::<()>::new("TEST");
        for op in &stream { apply(&book, op); }

        let bid_before = bid_qty_sum(&book);
        let ask_before = ask_qty_sum(&book);

        // Non-matching op: snapshot must not change resting qty.
        let _snap = book.snapshot();

        prop_assert_eq!(bid_qty_sum(&book), bid_before);
        prop_assert_eq!(ask_qty_sum(&book), ask_before);
    }
}