piker/tests/test_fsp_sync.py

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