piker/tests/test_trend_sampling.py

# piker: trading gear for hackers
# Copyright (C) 2018-present Tyler Goodlet (in stewardship of pikers)
# SPDX-License-Identifier: AGPL-3.0-or-later

# 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/>.

'''
Native sampler timestamp ordering required by causal processors.

'''
from collections import defaultdict
from collections.abc import Iterator
from types import SimpleNamespace
from unittest.mock import patch

import numpy as np
import trio
from trio.testing import MockClock

from piker.data._sampling import Sampler
from piker.fsp._trend import TrendCalculator, TrendHistory


class SampleSource:
    '''
    Provide the actual sampler's OHLC and append protocol in memory.

    '''
    def __init__(self) -> None:
        '''
        Seed one closed row and the sampler's future open tail.

        '''
        self._array: np.ndarray = np.zeros(16, dtype=[
            ('index', '<i8'),
            ('time', '<f8'),
            ('open', '<f8'),
            ('high', '<f8'),
            ('low', '<f8'),
            ('close', '<f8'),
            ('volume', '<f8'),
        ])
        self._array['index'] = np.arange(16)
        self._array['time'][:2] = [1000, 1001]
        self._array['close'][:2] = [3, 4]
        self._write_fields: list[str] = list(
            self._array.dtype.names[1:],
        )
        self.last: int = 2

    @property
    def array(self) -> np.ndarray:
        '''
        Expose only the published prefix, including its open tail.

        '''
        return self._array[:self.last]

    def push(self, values: np.ndarray) -> None:
        '''
        Append OHLC fields while retaining absolute source indices.

        '''
        end: int = self.last + len(values)
        self._array[self._write_fields][self.last:end] = values
        self.last = end


def test_sampler_waits_for_strictly_new_epoch() -> None:
    '''
    Empty realtime feeds seed times round(now) and round(now)+1.
    Their first sampler wake can have the same epoch as that future
    tail; rounded wall time can also repeat or move backward later.
    Appending those epochs makes strict causal FSP bootstrap fail.
    Drive the real sampler through all cases with a controlled clock:
    keep broadcasting wakeups, append only newer epochs, and verify
    the complete output remains valid Trend Calculus input.

    '''
    async def main() -> None:
        '''
        Advance virtual time without starting brokers or actors.

        '''
        source: SampleSource = SampleSource()
        epochs: tuple[int, ...] = (
            1001, 1001, 1000, 1002, 1002, 1003,
        )
        clock_values: Iterator[int] = iter(epochs)
        counts: list[int] = []
        complete: trio.Event = trio.Event()

        class ControlledSampler(Sampler):
            '''
            Isolate sampler state and capture every step broadcast.

            '''
            ohlcv_shms: dict = {1: [source]}
            subscribers: defaultdict = defaultdict(
                lambda: [0, set()],
            )

            @classmethod
            async def broadcast(cls, **kwargs: object) -> None:
                '''
                Stop after all scheduled clock cases are observed.

                '''
                counts.append(len(source.array))
                if len(counts) == len(epochs):
                    complete.set()
                    await trio.sleep_forever()

        fake_time: SimpleNamespace = SimpleNamespace(
            time=lambda: next(clock_values),
        )
        with patch('piker.data._sampling.time', fake_time):
            nursery: trio.Nursery
            async with trio.open_nursery() as nursery:
                scope: trio.CancelScope = await nursery.start(
                    ControlledSampler.increment_ohlc_buffer, 1.,
                )
                await complete.wait()
                scope.cancel()

        assert counts == [2, 2, 2, 3, 3, 4]
        assert source.array['time'].tolist() == [
            1000, 1001, 1002, 1003,
        ]
        result: TrendHistory = TrendCalculator().extend(
            source.array['close'],
            source.array['time'],
            source.array['index'],
        )
        assert len(result.trends) == 4

    trio.run(main, clock=MockClock(autojump_threshold=0))