piker/tests/providers/test_history_reservations.py

249 lines
9.0 KiB
Python

# 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)