tests/test_tick_factor.py

"""Tests for the tick-level trade-size + buy-initiated ratio factors."""

from __future__ import annotations

import numpy as np
import pandas as pd
import pytest

from scanner.factor_sources import StubDataSource
from scanner.tick_factor import (
    _bucket,
    _sign_trades,
    compute_tick_factors,
    compute_tick_factors_batch,
)


def test_bucket_boundaries():
    assert _bucket(50) == "retail"
    assert _bucket(99) == "retail"
    assert _bucket(100) == "small"
    assert _bucket(999) == "small"
    assert _bucket(1_000) == "medium"
    assert _bucket(9_999) == "medium"
    assert _bucket(10_000) == "block"
    assert _bucket(100_000) == "block"


def _ticks(prices, sizes, bid=99.99, ask=100.01):
    """Build a tick frame with the given prices/sizes at a fixed mid."""
    n = len(prices)
    return pd.DataFrame({
        "ts": pd.date_range("2026-06-02 09:30", periods=n, freq="1s"),
        "price": prices,
        "size": sizes,
        "bid": [bid] * n,
        "ask": [ask] * n,
    })


def test_sign_at_ask_is_buy():
    df = _ticks(prices=[100.02, 100.02], sizes=[100, 100])
    signs = _sign_trades(df)
    assert (signs == 1).all()


def test_sign_at_bid_is_sell():
    df = _ticks(prices=[99.98, 99.98], sizes=[100, 100])
    signs = _sign_trades(df)
    assert (signs == -1).all()


def test_sign_at_mid_carries_forward():
    df = _ticks(
        prices=[100.02, 100.00, 100.00, 99.98],  # buy, mid, mid, sell
        sizes=[100, 100, 100, 100],
    )
    signs = _sign_trades(df).tolist()
    # [1, carry(1), carry(1), -1]
    assert signs == [1, 1, 1, -1]


def test_block_share_calculation():
    df = _ticks(
        prices=[100.02] * 4,
        sizes=[100, 1000, 5000, 15000],   # 15k block out of 21.1k total
    )
    f = compute_tick_factors("X", source=_StaticTickSource(df))
    # 15000 / (100 + 1000 + 5000 + 15000) = 0.71
    assert 0.70 < f["block_share"] < 0.72


def test_block_aggression_positive_when_block_buys():
    df = _ticks(
        prices=[100.02] * 3,   # all buys (at ask)
        sizes=[500, 5000, 20000],  # 20k block
    )
    f = compute_tick_factors("X", source=_StaticTickSource(df))
    assert f["block_aggression"] == pytest.approx(1.0)


def test_block_aggression_negative_when_block_sells():
    df = _ticks(
        prices=[99.98] * 3,   # all sells (at bid)
        sizes=[500, 5000, 20000],
    )
    f = compute_tick_factors("X", source=_StaticTickSource(df))
    assert f["block_aggression"] == pytest.approx(-1.0)


def test_buy_ratio_in_unit_interval():
    df = _ticks(
        prices=[100.02, 100.02, 99.98, 99.98],
        sizes=[100, 200, 300, 400],
    )
    f = compute_tick_factors("X", source=_StaticTickSource(df))
    assert 0.0 <= f["buy_ratio"] <= 1.0


def test_empty_ticks():
    f = compute_tick_factors("X", source=_StaticTickSource(pd.DataFrame()))
    assert f["block_share"] == 0.0
    assert f["block_aggression"] == 0.0
    assert f["buy_ratio"] == 0.5


def test_no_bid_ask_falls_back_to_rolling_mid():
    df = pd.DataFrame({
        "ts": pd.date_range("2026-06-02 09:30", periods=50, freq="1s"),
        "price": 100 + np.cumsum(np.random.default_rng(1).normal(0, 0.01, 50)),
        "size": [100] * 50,
    })
    f = compute_tick_factors("X", source=_StaticTickSource(df))
    assert "buy_ratio" in f


def test_batch():
    df = _ticks(prices=[100.02] * 2, sizes=[500, 10000])
    src = _StaticTickSource(df)
    out = compute_tick_factors_batch(["A", "B"], source=src)
    assert set(out.index) == {"A", "B"}
    assert "block_share" in out.columns


def test_stub_synthesises():
    src = StubDataSource(stub_dir="/nonexistent")
    df = src.get_ticks("AAPL")
    assert df is not None and not df.empty
    assert {"ts", "price", "size"}.issubset(df.columns)


# --- helper source ---
class _StaticTickSource:
    def __init__(self, df):
        self._df = df
    def get_ticks(self, ticker, date=None):
        return self._df