piker/tests/test_history_backfill.py

347 lines
9.8 KiB
Python

'''
Deterministic history-backfill regressions.
'''
from functools import partial
from pathlib import Path
from types import SimpleNamespace
from uuid import uuid4
import numpy as np
from pendulum import (
datetime,
from_timestamp,
)
import polars as pl
import pytest
import tractor
import trio
from piker.brokers import DataUnavailable
from piker.brokers.ib.api import (
_bar_load_dtype,
_ohlc_dtype,
)
from piker.data._sharedmem import maybe_open_shm_array
from piker.data._source import def_iohlcv_fields
from piker.storage.nativedb import (
NativeStorageClient,
ohlc_key_map,
)
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_ib_latest_frame_round_trips_through_nativedb(
tmp_path: Path,
) -> None:
'''
Persist and reload IB's provider schema through history startup.
``publish_latest_frame()`` previously passed IB's index-less bars
directly to NativeDB, whose durable-schema validator rejected the
missing derived ``index``. IB also appends ``count`` after OHLCV,
which exposed positional Polars-to-NumPy conversion to field
corruption. Build the actual IB load dtype with distinct values,
run the real history publication, NativeDB Parquet, and ``ShmArray``
paths, then hydrate a fresh IB buffer from storage. Assertions prove
first-start publication preserves provider fields while restart maps
canonical fields and leaves provider-only ``count`` at its default.
'''
frame = np.zeros(
2,
dtype=np.dtype(_bar_load_dtype),
)
frame['time'] = [60, 120]
frame['open'] = [1.1, 2.1]
frame['high'] = [1.2, 2.2]
frame['low'] = [1.0, 2.0]
frame['close'] = [1.15, 2.15]
frame['volume'] = [10, 20]
frame['count'] = [3, 4]
storage = NativeStorageClient(tmp_path)
shm_key = f'test_ib_history_{uuid4().hex}'
mkt = SimpleNamespace(
fqme='mnq.cme.20260918.ib',
dst=SimpleNamespace(atype='continuous_future'),
src=SimpleNamespace(atype='fiat'),
get_fqme=lambda **kwargs: 'mnq.cme.20260918.ib',
)
async def main() -> None:
with trio.fail_after(2):
async with tractor.open_root_actor(
name=shm_key,
tpt_bind_addrs=[('127.0.0.1', 0)],
):
shm, opened = maybe_open_shm_array(
key=f'{shm_key}_first',
size=16,
dtype=np.dtype(_ohlc_dtype),
append_start_index=8,
)
restart_shm, restart_opened = maybe_open_shm_array(
key=f'{shm_key}_restart',
size=16,
dtype=np.dtype(_ohlc_dtype),
append_start_index=8,
)
assert opened
assert restart_opened
await publish_latest_frame(
storage=storage,
mkt=mkt,
shm=shm,
array=frame,
timeframe=60,
)
loaded = await storage.read_ohlcv(
mkt.fqme,
timeframe=60,
)
restart_shm.push(
loaded,
prepend=True,
field_map=ohlc_key_map,
)
stored = pl.read_parquet(
storage.mk_path(mkt.fqme, 60)
)
canonical = [
name
for name, _ in def_iohlcv_fields
]
canonical_schema = {
name: (
pl.Int64
if field_type is int
else pl.Float64
)
for name, field_type in def_iohlcv_fields
}
assert stored.columns == canonical
assert dict(stored.schema) == canonical_schema
assert list(loaded.dtype.fields) == canonical
assert loaded['index'].tolist() == [0, 1]
for field in canonical[1:]:
assert (
loaded[field].tolist()
==
frame[field].tolist()
)
assert (
restart_shm.array[field].tolist()
==
frame[field].tolist()
)
assert shm.array['count'].tolist() == [3, 4]
assert restart_shm.array['count'].tolist() == [0, 0]
trio.run(main)
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)