piker/tests/test_fsp_sync.py

134 lines
3.6 KiB
Python

'''
FSP history synchronization regressions.
'''
from collections.abc import AsyncIterator
from typing import cast
import numpy as np
import pytest
import trio
from piker.fsp._momo import (
rsi,
wma,
)
from piker.fsp._volume import (
tina_vwap,
)
from piker.data._sharedmem import NDTokenMsg
from piker.data.ticktools import FeedQuote
from piker.fsp._api import Fsp
from tractor.ipc._shm import (
NDToken,
ShmArray,
)
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,
)
@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)