piker/tests/test_history_backfill.py

'''
Deterministic history-backfill regressions.

'''
from functools import partial
from types import SimpleNamespace

import numpy as np
from pendulum import (
    datetime,
    from_timestamp,
)
import pytest
import trio

from piker.brokers import DataUnavailable
from piker.data._source import def_iohlcv_fields
from piker.tsp._history import (
    notify_backfill,
    publish_latest_frame,
    start_backfill,
)


def run_exhausted_backfill(get_hist) -> None:
    '''
    Run a provider-exhaustion case without actor services.

    '''
    async def main() -> None:
        with trio.fail_after(0.5):
            async with trio.open_nursery() as nursery:
                await nursery.start(partial(
                    start_backfill,
                    get_hist=get_hist,
                    def_frame_duration=None,
                    mod=SimpleNamespace(name='fake'),
                    mkt=SimpleNamespace(fqme='x.fake'),
                    shm=object(),
                    timeframe=60,
                    backfill_from_shm_index=100,
                    backfill_from_dt=datetime(2026, 1, 2),
                    sampler_stream=object(),
                    backfill_until_dt=datetime(2026, 1, 1),
                    storage=None,
                    write_tsdb=False,
                ))

    trio.run(main)


def test_data_unavailable_completes() -> None:
    '''
    Provider exhaustion exits without orphaned completion waits.

    '''
    async def get_hist(*args, **kwargs):
        raise DataUnavailable('history exhausted')

    run_exhausted_backfill(get_hist)


def test_empty_frame_completes() -> None:
    '''
    An empty provider frame is checked before endpoint indexing.

    '''
    frame = np.empty(
        0,
        dtype=np.dtype(def_iohlcv_fields),
    )

    async def get_hist(*args, **kwargs):
        end_dt = kwargs['end_dt']
        return frame, end_dt, end_dt

    run_exhausted_backfill(get_hist)


def test_storage_receives_full_provider_delta() -> None:
    '''
    SHM capacity truncation does not truncate durable history.

    '''
    frame = np.zeros(
        2,
        dtype=np.dtype(def_iohlcv_fields),
    )
    frame['index'] = [0, 1]
    frame['time'] = [60, 120]

    events: list[str] = []

    class Shm:
        def __init__(self) -> None:
            self.pushed: list[np.ndarray] = []

        def push(self, array, **kwargs) -> None:
            self.pushed.append(array.copy())
            events.append('shm')

    class Sampler:
        async def send(self, msg) -> None:
            events.append('sampler')

    class Storage:
        def __init__(self) -> None:
            self.frames: list[np.ndarray] = []

        async def update_ohlcv(
            self,
            fqme: str,
            ohlcv: np.ndarray,
            timeframe: int,
        ) -> None:
            self.frames.append(ohlcv.copy())
            events.append('storage')

    shm = Shm()
    storage = Storage()
    mkt = SimpleNamespace(
        fqme='x.test',
        dst=SimpleNamespace(atype='crypto'),
        src=SimpleNamespace(atype='crypto_currency'),
        get_fqme=lambda **kwargs: 'x.test',
    )

    async def get_hist(*args, **kwargs):
        return (
            frame,
            from_timestamp(60),
            from_timestamp(120),
        )

    async def main() -> None:
        with trio.fail_after(0.5):
            await start_backfill(
                get_hist=get_hist,
                def_frame_duration=None,
                mod=SimpleNamespace(name='fake'),
                mkt=mkt,
                shm=shm,
                timeframe=60,
                backfill_from_shm_index=1,
                backfill_from_dt=from_timestamp(180),
                sampler_stream=Sampler(),
                backfill_until_dt=from_timestamp(60),
                storage=storage,
                write_tsdb=True,
            )

    trio.run(main)
    assert shm.pushed[0]['time'].tolist() == [120]
    assert storage.frames[0]['time'].tolist() == [60, 120]
    assert events == ['storage', 'shm', 'sampler']


def test_latest_frame_is_persisted_before_shm() -> None:
    '''
    Startup persists the most-recent provider frame before publication.

    '''
    frame = np.zeros(
        2,
        dtype=np.dtype(def_iohlcv_fields),
    )
    frame['time'] = [60, 120]
    events: list[str] = []

    class Storage:
        async def update_ohlcv(self, *args) -> None:
            events.append('storage')

    class Shm:
        def push(self, array, **kwargs) -> None:
            events.append('shm')

    async def main() -> None:
        with trio.fail_after(0.5):
            await publish_latest_frame(
                storage=Storage(),
                mkt=SimpleNamespace(
                    fqme='x.test',
                    dst=SimpleNamespace(atype='crypto'),
                    src=SimpleNamespace(atype='crypto_currency'),
                    get_fqme=lambda **kwargs: 'x.test',
                ),
                shm=Shm(),
                array=frame,
                timeframe=60,
            )

    trio.run(main)
    assert events == ['storage', 'shm']


def test_backfill_notification_timeout_is_bounded(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    '''
    A wedged sampler can not indefinitely shield actor teardown.

    '''
    monkeypatch.setattr(
        'piker.tsp._history._notify_timeout_s',
        0.01,
    )

    class Sampler:
        async def send(self, msg) -> None:
            await trio.sleep_forever()

    async def main() -> None:
        with trio.fail_after(0.1):
            await notify_backfill(
                Sampler(),
                SimpleNamespace(fqme='x.test'),
                60,
            )

    trio.run(main)