piker/tests/providers/test_history_reservations.py

# piker: trading gear for hackers
# Copyright (C) 2018-present  Tyler Goodlet (in stewardship of pikers)

# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.

# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.

# You should have received a copy of the GNU Affero General Public License
# along with this program.  If not, see <https://www.gnu.org/licenses/>.

'''
Schedule-driven startup allocation without published-row movement.

'''
from functools import partial
from pathlib import Path

import numpy as np
import pytest
import trio

from piker.data._source import def_iohlcv_fields
from piker.tsp._reserve import reserve_history_gaps


@pytest.mark.parametrize('case', [
    'covered',  # Reserve open-session slots; exclude the closure.
    'partial',  # Coverage misses the gap's left margin.
    'timeout',  # Provider failure is not evidence of a closure.
    'rth',  # Regular-hours sessions omit extended trading hours.
    'capacity',  # Too many placeholders: leave the gap unresolved.
    'empty',  # No sessions: do not assume the venue was closed.
])
def test_reservation_requires_covered_extended_sessions(case):
    '''
    Interrupted backfill leaves packed rows around missing time.

    Model a closure inside that interval: reserve only open-session
    samples and retain explicit synthetic ranges. Partial coverage,
    provider failure, RTH-only replies, and insufficient capacity
    must never invent a schedule. Input rows remain immutable and
    output placeholders carry zero activity, not invented trades.

    '''
    frame = np.zeros(2, dtype=np.dtype(def_iohlcv_fields))
    frame['time'] = [100, 110]
    frame['close'] = [7, 9]
    frame['volume'] = [4, 5]
    before = frame.copy()
    async def query(left, right):
        assert (left, right) == (100, 110)
        if case == 'timeout':
            raise TimeoutError
        return {
            'start': 102 if case == 'partial' else 100,
            'end': 111,
            'sessions': [] if case == 'empty' else [
                (100, 104), (108, 111),
            ],
            'use_rth': case == 'rth',
        }
    result, ranges = trio.run(partial(
        reserve_history_gaps, frame, 1, query,
        max_rows=3 if case == 'capacity' else 20,
    ))
    np.testing.assert_array_equal(frame, before)
    if case == 'covered':
        assert result['time'].tolist() == [100, 101, 102, 103,
                                           108, 109, 110]
        assert ranges == [(101, 104), (108, 110)]
        assert result['volume'].tolist() == [4, 0, 0, 0, 0, 0, 5]
        assert result['close'].tolist() == [7, 7, 7, 7, 7, 7, 9]
    else:
        np.testing.assert_array_equal(result, before)
        assert ranges == []


@pytest.mark.parametrize('newest_gaps', [-1, 0, 1, 2])
def test_reservation_limits_newest_gap_requests(
    newest_gaps: int,
) -> None:
    '''
    A configurable request budget must bound provider calls.

    Two gaps distinguish newest-first selection from array order.
    Zero must disable queries, not select all gaps via `[-0:]`;
    negative budgets are rejected before any provider call. Check
    both fetched intervals and resulting synthetic timestamps.

    '''
    frame: np.ndarray = np.zeros(3, dtype=def_iohlcv_fields)
    frame['time'] = [100, 103, 106]
    calls: list[tuple[float, float]] = []

    async def query(left: float, right: float) -> dict:
        calls.append((left, right))
        return {
            'start': left,
            'end': right,
            'use_rth': False,
            'sessions': [(left, right)],
        }

    run = partial(
        reserve_history_gaps,
        frame, 1, query, max_rows=20, newest_gaps=newest_gaps,
    )
    if newest_gaps < 0:
        with pytest.raises(ValueError, match='nonnegative'):
            trio.run(run)
        assert calls == []
        return

    result, ranges = trio.run(run)
    assert calls == [(103, 106), (100, 103)][:newest_gaps]
    expected: dict[int, list[int]] = {
        0: [100, 103, 106],
        1: [100, 103, 104, 105, 106],
        2: [100, 101, 102, 103, 104, 105, 106],
    }
    assert result['time'].tolist() == expected[newest_gaps]
    assert len(ranges) == newest_gaps
    np.testing.assert_array_equal(frame['time'], [100, 103, 106])


def test_startup_reserves_shm_without_persisting_synthetic_rows(
    tmp_path: Path,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    '''
    Restart after interrupted backfill must retain gap provenance.

    Seed parquet with two separated observations, then run the real
    startup backfiller and SHM prepend. A fake schedule supplies a
    maintenance closure inside the hole. Wait for the actual tail
    repair to complete before cancelling the service task. Assert
    reserved open-session rows exist in SHM with queryable metadata,
    while a fresh parquet load contains only provider observations.

    '''
    from contextlib import asynccontextmanager
    from types import SimpleNamespace
    from uuid import uuid4

    from pendulum import duration, from_timestamp
    import tractor

    from piker.data._sharedmem import maybe_open_shm_array
    from piker.storage.nativedb import (
        NativeStorageClient, ohlc_key_map,
    )
    from piker.tsp import _history
    from piker.tsp._reserve import (
        reservation_ranges,
    )

    epoch = 1790193489
    fqme = 'schedule.test'
    finished = trio.Event()
    def bars(offsets):
        frame = np.zeros(len(offsets), dtype=def_iohlcv_fields)
        frame['time'] = np.asarray(offsets) + epoch
        for field in ('open', 'high', 'low', 'close'):
            frame[field] = 100
        frame['volume'] = 1
        return frame
    async def history(tf, end_dt=None, **kw):
        frame = bars([11, 12] if end_dt is None else [10, 11])
        return (frame, from_timestamp(frame['time'][0]),
                from_timestamp(frame['time'][-1]))
    async def schedule(left, right):
        return {
            'start': epoch, 'end': epoch + 13, 'use_rth': False,
            'sessions': [
                (epoch, epoch + 4), (epoch + 8, epoch + 13),
            ],
        }
    @asynccontextmanager
    async def history_client(*args, **kwargs):
        yield history, {
            'frame_types': {1: duration(seconds=2)},
            'query_schedule': schedule,
        }
    real_repair = _history.maybe_fill_null_segments
    async def repair(**kwargs):
        await real_repair(**kwargs)
        finished.set()
    monkeypatch.setattr(_history, 'maybe_fill_null_segments', repair)
    # Isolate actor-local provenance as well as SHM and parquet.
    monkeypatch.setattr(
        'piker.tsp._reserve._reservations', {},
    )
    class Sampler:
        async def send(self, msg):
            pass
    async def main():
        store = NativeStorageClient(tmp_path)
        await store.update_ohlcv(fqme, bars([0, 10]), 1)
        with trio.fail_after(5):
            async with tractor.open_root_actor(
                name='schedule-test',
                tpt_bind_addrs=[('127.0.0.1', 0)],
            ):
                shm, opened = maybe_open_shm_array(
                    key=f'schedule_{uuid4().hex}', size=100,
                    dtype=np.dtype(def_iohlcv_fields),
                    append_start_index=80,
                )
                assert opened
                async with trio.open_nursery() as nursery:
                    nursery.start_soon(partial(
                        _history.tsdb_backfill,
                        mod=SimpleNamespace(
                            name='fake',
                            open_history_client=history_client,
                        ),
                        storemod=SimpleNamespace(
                            ohlc_key_map=ohlc_key_map,
                        ),
                        storage=store,
                        mkt=SimpleNamespace(
                            fqme=fqme, get_fqme=lambda **kw: fqme,
                            dst=SimpleNamespace(atype='future'),
                            src=SimpleNamespace(atype='fiat'),
                        ),
                        shm=shm, timeframe=1,
                        sampler_stream=Sampler(),
                    ))
                    await finished.wait()
                    nursery.cancel_scope.cancel()
                assert (shm.array['time'] - epoch).tolist() == [
                    0, 1, 2, 3, 8, 9, 10, 11, 12,
                ]
                assert reservation_ranges(fqme, 1) == [
                    (epoch + 1, epoch + 4), (epoch + 8, epoch + 10),
                ]
                fresh_store = NativeStorageClient(tmp_path)
                disk, _, _ = await fresh_store.load(fqme, 1)
                assert (disk['time'] - epoch).tolist() == [
                    0, 10, 11, 12,
                ]
    trio.run(main)