piker/tests/test_fsp_sync.py

'''
FSP history synchronization regressions.

'''
from collections.abc import AsyncIterator
from types import SimpleNamespace
from typing import cast

import numpy as np
import pytest
import trio

from piker.fsp._engine import (
    _is_relevant_sample_msg,
    _needs_history_resync,
    _open_history_snapshot,
    _should_advance_dst,
    Cascade,
    EdgeFunc,
    QuoteFrame,
)
from piker.fsp._momo import (
    rsi,
    wma,
)
from piker.fsp._volume import (
    dolla_vlm,
    tina_vwap,
)
from piker.accounting import MktPair
from piker.data.flows import Flume
from piker.data._sharedmem import NDTokenMsg
from piker.data.ticktools import FeedQuote
from piker.fsp._api import Fsp
from piker.ui._chart import LinkedSplits
from piker.ui._fsp import update_fsp_vizs
from tractor.ipc._shm import (
    NDToken,
    ShmArray,
)

type SampleMsg = dict[
    str,
    int|float|tuple[str, float],
]


def test_sample_period_ignores_market_closure_gap() -> None:
    '''
    Keep FSP cascades subscribed to their regular sampler period.

    Starting a cascade immediately after a market closure left the
    source SHM tail with a large gap between its last two distinct
    timestamps. Using only that pair subscribed the cascade to the gap
    duration instead of the regular one-second stream. Realtime FSP
    writes then repeatedly replaced one row without advancing its SHM
    bound. Give both source SHMs closure-gap tail deltas and prove
    `Flume.get_ds_info()` returns its declared one- and 60-second
    periods instead of deriving cadence from those rows.

    '''
    rt_shm = OhlcvShm(length=4)
    hist_shm = OhlcvShm(length=4)
    rt_shm._array['time'][:4] = [100, 101, 102, 3600]
    hist_shm._array['time'][:4] = [100, 160, 220, 3600]
    flume = Flume(
        mkt=cast(MktPair, SimpleNamespace()),
        first_quote={},
        _rt_shm_token=rt_shm._token,
        _hist_shm_token=hist_shm._token,
    )
    flume._rt_shm = cast(ShmArray, rt_shm)
    flume._hist_shm = cast(ShmArray, hist_shm)

    assert np.diff(rt_shm.array['time'])[-1] == 3498
    assert np.diff(hist_shm.array['time'])[-1] == 3380
    assert flume.get_ds_info() == (1, 60, 60)


class Value:
    def __init__(self, value: int) -> None:
        self.value: int = value


class Shm:
    def __init__(
        self,
        first: int,
        last: int,
        token: str = 'fsp',

    ) -> None:
        self._first: Value = Value(first)
        self._last: Value = Value(last)
        self._array: np.ndarray = np.ones(4096)
        self._len: int = len(self._array)
        self._token: NDToken = NDToken(
            shm_name=token,
            shm_first_index_name=f'{token}_first',
            shm_last_index_name=f'{token}_last',
            dtype_descr=(('value', '<f8'),),
            size=len(self._array),
        )

    @property
    def array(self) -> np.ndarray:
        return self._array[
            self._first.value:self._last.value
        ]

    @property
    def token(self) -> NDTokenMsg:
        return cast(NDTokenMsg, self._token.as_msg())

    @property
    def index(self) -> int:
        return self._last.value % len(self._array)

    def last(self, length: int = 1) -> np.ndarray:
        return self.array[-length:]


class OhlcvShm(Shm):
    def __init__(self, length: int = 32) -> None:
        dtype = np.dtype([
            ('index', '<i8'),
            ('time', '<i8'),
            ('open', '<f8'),
            ('high', '<f8'),
            ('low', '<f8'),
            ('close', '<f8'),
            ('volume', '<f8'),
        ])
        self._first = Value(0)
        self._last = Value(length)
        self._array = np.ones(length + 8, dtype=dtype)
        self._array['index'] = np.arange(length + 8)
        self._array['time'] = np.arange(length + 8)
        self._array['close'] = np.arange(length + 8) + 1
        self._array['high'] = self._array['close'] + 1
        self._array['low'] = self._array['close'] - 1
        self._len = len(self._array)
        self._token = NDToken(
            shm_name='ohlcv',
            shm_first_index_name='ohlcv_first',
            shm_last_index_name='ohlcv_last',
            dtype_descr=tuple(dtype.descr),
            size=self._len,
        )


class FspShm(Shm):
    def __init__(
        self,
        first: int = 0,
        last: int = 4,
        token: str = 'fsp',

    ) -> None:
        dtype = np.dtype([
            ('index', '<i8'),
            ('time', '<f8'),
            ('flow', '<f8'),
            ('dark_flow', '<f8'),
        ])
        size: int = max(last + 2, 4096)
        self._first = Value(first)
        self._last = Value(last)
        self._array = np.zeros(size, dtype=dtype)
        self._array['index'] = np.arange(size)
        self._array['time'] = np.arange(size)
        self._array['flow'] = np.arange(size)
        self._len = len(self._array)
        self._token = NDToken(
            shm_name=token,
            shm_first_index_name=f'{token}_first',
            shm_last_index_name=f'{token}_last',
            dtype_descr=tuple(dtype.descr),
            size=self._len,
        )


def mk_cascade(
    src_bounds: tuple[int, int],
    dst_bounds: tuple[int, int],

) -> Cascade:
    fsp: Fsp = cast(
        Fsp,
        SimpleNamespace(
            name='test_fsp',
            ns_path='tests:test_fsp',
        ),
    )
    src: Flume = cast(
        Flume,
        SimpleNamespace(
            rt_shm=Shm(*src_bounds, token='src'),
        ),
    )
    dst: Flume = cast(
        Flume,
        SimpleNamespace(
            rt_shm=Shm(*dst_bounds, token='dst'),
        ),
    )
    nursery: trio.Nursery = cast(trio.Nursery, None)
    return Cascade(src, dst, nursery, fsp)


@pytest.mark.parametrize(
    'src_bounds,dst_bounds,expected',
    [
        ((10, 20), (10, 20), (True, 0, 0)),
        ((10, 21), (10, 20), (True, 1, 1)),
        ((0, 3000), (2000, 3000), (False, 0, 2000)),
        ((9, 20), (10, 20), (False, 0, 1)),
        ((8, 20), (10, 20), (False, 0, 2)),
        ((10, 20), (9, 20), (False, 0, 1)),
        ((11, 21), (10, 20), (False, 1, 0)),
        ((10, 22), (10, 20), (False, 2, 2)),
        ((10, 19), (10, 20), (False, -1, 1)),
    ],
)
def test_cascade_sync_uses_absolute_bounds(
    src_bounds: tuple[int, int],
    dst_bounds: tuple[int, int],
    expected: tuple[bool, int, int],

) -> None:
    '''
    Detect every history-bound mismatch without modulo-index skew.

    Source prepends move only its first bound, while a normal realtime
    step moves only its last bound. The former requires a historical
    recompute even for one or two rows; the latter permits exactly one
    destination append. This matrix also shifts equal-length bounds and
    puts the destination ahead to prove neither state is accepted.

    '''
    cascade = mk_cascade(src_bounds, dst_bounds)

    assert cascade.is_synced() == expected


@pytest.mark.parametrize(
    'msg,expected',
    [
        ({'index': 1}, True),
        ({'backfilling': ('tsla.nasdaq.ib', 1)}, True),
        ({'backfilling': ('btcusdt.binance', 1)}, False),
        ({'backfilling': ('tsla.nasdaq.ib', 60)}, False),
    ],
)
def test_cascade_filters_foreign_backfill_events(
    msg: SampleMsg,
    expected: bool,

) -> None:
    '''
    Ignore backfill wakeups from overlay markets and other periods.

    Samplerd broadcasts every backfill marker to every period and FSP
    subscriber. With multiple markets overlaid, an unrelated provider
    frame used to wake each primary-market cascade and race its own SHM
    inspection. Feed regular sample steps and matching history events
    through, but reject both a foreign FQME and the 60-second period.

    '''
    assert _is_relevant_sample_msg(
        msg,
        fqme='tsla.nasdaq.ib',
        period_s=1,
    ) is expected


@pytest.mark.parametrize(
    'msg,step_diff,expected',
    [
        ({'index': 1}, 1, True),
        ({'index': 1}, 0, False),
        ({'backfilling': ('tsla.nasdaq.ib', 1)}, 1, False),
        ({'backfilling': ('tsla.nasdaq.ib', 1)}, 0, False),
    ],
)
def test_only_new_sample_steps_advance_destination(
    msg: SampleMsg,
    step_diff: int,
    expected: bool,

) -> None:
    '''
    Keep queued history and duplicate wakeups from adding FSP rows.

    A recomputation can consume several prepends before their sampler
    markers are delivered. Those stale markers then observe aligned SHM
    bounds and previously appended duplicate destination rows. Model the
    queued marker and regular duplicate cases, and prove only a genuine
    one-step source lead authorizes destination advancement.

    '''
    assert _should_advance_dst(msg, step_diff) is expected


@pytest.mark.parametrize(
    'msg,synced,expected',
    [
        ({'index': 1}, True, False),
        ({'index': 1}, False, True),
        ({'backfilling': ('tsla.nasdaq.ib', 1)}, True, True),
        ({'backfilling': ('tsla.nasdaq.ib', 1)}, False, True),
    ],
)
def test_backfill_always_invalidates_fsp_history(
    msg: SampleMsg,
    synced: bool,
    expected: bool,

) -> None:
    '''
    Recompute in-place history repairs with unchanged SHM bounds.

    Gap repair can replace null source rows without moving either SHM
    bound. A bounds-only predicate therefore reports synchronization
    even though derived values are stale. Prove every relevant backfill
    marker forces history replay while an aligned sample event does not.

    '''
    assert _needs_history_resync(msg, synced) is expected


def test_history_compute_uses_immutable_source_snapshot() -> None:
    '''
    Do not publish output and timestamps from different source ranges.

    A source append or prepend can occur while the FSP's first yield is
    computing. The old path read timestamps only afterward and could pair
    an N-row result with a shifted N-row timestamp tail. Mutate the live
    last bound during the first yield, then prove the accepted output and
    timestamps retain the original range while later reads switch to live
    SHM.

    '''
    shm = Shm(10, 20, token='source')
    live_lengths: list[int] = []

    async def edge(
        _source: AsyncIterator[FeedQuote],
        src_shm: ShmArray,

    ) -> AsyncIterator[np.ndarray]:
        shm._last.value += 1
        yield np.ones(len(src_shm.array))
        live_lengths.append(len(src_shm.array))

    async def main() -> None:
        (
            out_stream,
            history,
            bounds,
            src_history,
        ) = await _open_history_snapshot(
            cast(EdgeFunc, edge),
            quote_stream=cast(
                AsyncIterator[QuoteFrame],
                object(),
            ),
            fqme='tsla.nasdaq.ib',
            src_shm=cast(ShmArray, shm),
        )
        assert bounds == (10, 20)
        assert len(history) == 10
        assert len(src_history) == 10
        with pytest.raises(StopAsyncIteration):
            await anext(out_stream)

    trio.run(main)

    assert live_lengths == [11]


@pytest.mark.parametrize('target', [wma, rsi, tina_vwap])
def test_builtin_fsp_stream_contract(target: Fsp) -> None:
    '''
    Keep every registered scalar FSP on the engine's yield protocol.

    The momentum operators previously had incompatible call signatures,
    short historical arrays, and bare realtime yields, all hidden by an
    engine-side cast. Run each against one OHLCV snapshot and one trade,
    proving the first yield is a source-aligned array and the next yield
    is a named realtime field/value pair.

    '''
    shm = cast(ShmArray, OhlcvShm())

    async def source() -> AsyncIterator[FeedQuote]:
        yield {
            'ticks': [{
                'type': 'trade',
                'price': 42.0,
                'size': 1.0,
            }],
        }

    async def main() -> None:
        stream = target.func(source(), shm)
        history = await anext(stream)
        assert isinstance(history, np.ndarray)
        assert len(history) == len(shm.array)

        realtime = await anext(stream)
        assert isinstance(realtime, tuple)
        assert realtime[0] == target.name
        await stream.aclose()

    trio.run(main)


def test_dolla_vlm_rejects_negative_trade_volume() -> None:
    '''
    Keep a provider volume reset from hiding the zoomed-out curve.

    IB can emit a large negative trade size around a session closure.
    The sampler previously stored that size as bar volume and the
    dollar-volume history multiplied it by price. The resulting outlier
    made `Viz.maxmin()` reject the entire downsampled range, leaving a
    stale cached path or no visible curve. Seed one historical negative
    volume and stream one negative trade tick, then prove both historical
    and realtime dollar volume remain nonnegative.

    '''
    shm = OhlcvShm(length=4)
    shm._array['volume'][:4] = [1, 2, -2_319_096, 3]

    async def source() -> AsyncIterator[FeedQuote]:
        yield {
            'ticks': [{
                'type': 'trade',
                'price': 29_495.75,
                'size': -2_319_096.,
            }],
        }

    async def main() -> None:
        stream = dolla_vlm.func(
            source(),
            cast(ShmArray, shm),
        )
        history = await anext(stream)
        assert isinstance(history, dict)
        values = history['dolla_vlm']
        assert isinstance(values, np.ndarray)
        assert values[2] == 0
        assert np.all(values >= 0)

        realtime = await anext(stream)
        assert realtime == ('dolla_vlm', 0)
        await stream.aclose()

    trio.run(main)


class Viz:
    def __init__(self, name: str, shm: Shm) -> None:
        self.name: str = name
        self.shm: Shm = shm
        self.updates: int = 0
        self.force_redraws: list[bool] = []
        self._last_fsp_update_sig: (
            tuple[int, int, bytes]|None
        ) = None
        self._mxmns: dict[
            tuple[int, int],
            tuple[int, int],
        ] = {(10, 20): (1, 2)}
        self.view = SimpleNamespace(
            rescale_count=0,
        )

        def rescale(
            *,
            do_linked_charts: bool,
            do_overlay_scaling: bool,

        ) -> None:
            assert not do_linked_charts
            assert do_overlay_scaling
            self.view.rescale_count += 1

        self.view.interact_graphics_cycle = rescale
        self.plot: SimpleNamespace = SimpleNamespace(
            getAxis=lambda name: SimpleNamespace(_stickies={}),
            vb=self.view,
        )

    def update_graphics(
        self,
        force_redraw: bool = False,

    ) -> None:
        self.updates += 1
        self.force_redraws.append(force_redraw)


def test_fsp_history_update_redraws_only_derived_vizs() -> None:
    '''
    Keep a source-history prepend from refreshing the whole chart.

    FSP cascades recompute after each near-term backfill frame. The old
    notification called the linked chart's full graphics cycle, which
    redrew the primary market and every overlay and disturbed the live
    view repeatedly. Arrange two visualizations sharing the rebuilt FSP
    SHM plus unrelated source and overlay SHMs, then prove only the two
    derived curves receive an update.

    '''
    fsp_shm = FspShm(10, 20, token='derived')
    source_viz = Viz('source', Shm(10, 20, token='source'))
    first_fsp_viz = Viz('flow', fsp_shm)
    second_fsp_viz = Viz('dark_flow', fsp_shm)
    overlay_viz = Viz('overlay', Shm(10, 20, token='overlay'))

    linked: LinkedSplits = cast(
        LinkedSplits,
        SimpleNamespace(
            chart=SimpleNamespace(
                _vizs={
                    'source': source_viz,
                    'overlay': overlay_viz,
                },
            ),
            subplots={
                'volume': SimpleNamespace(
                    _vizs={
                        'flow': first_fsp_viz,
                        'dark_flow': second_fsp_viz,
                    },
                ),
            },
        ),
    )

    updated = update_fsp_vizs(linked, fsp_shm._token)

    assert updated == 2
    assert first_fsp_viz.updates == 1
    assert second_fsp_viz.updates == 1
    assert first_fsp_viz.force_redraws == [True]
    assert second_fsp_viz.force_redraws == [True]
    assert first_fsp_viz._mxmns == {}
    assert second_fsp_viz._mxmns == {}
    assert first_fsp_viz.view.rescale_count == 1
    assert second_fsp_viz.view.rescale_count == 1
    assert source_viz.updates == 0
    assert overlay_viz.updates == 0