piker/tests/test_regime_feed.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 SHM-to-regime contracts, using synthetic source publications.

'''
from collections.abc import AsyncIterator, Callable
from contextlib import asynccontextmanager
from types import SimpleNamespace
from typing import Any, Never
import json

import httpx
import numpy as np
import pytest
import trio
from tractor.trionics import Lagged
from trio.testing import wait_all_tasks_blocked

from piker.fsp._regime import (
    CloseSample,
    MultiTimeframeTrend,
    RegimeSnapshot,
)
from piker.fsp._regime_feed import (
    ClosedSourceRows,
    SourceChanged,
    consume_regime_source,
    copy_closed_rows,
    run_regime_feed,
    serve_regimes,
)


class FakeSource:
    '''
    Model absolute bounds independently of mutable backing storage.

    '''
    def __init__(
        self,
        first: int = 0,
        last: int = 6,
    ) -> None:
        '''
        Seed indexed rows independently of their published bounds.

        '''
        self._first: SimpleNamespace = SimpleNamespace(value=first)
        self._last: SimpleNamespace = SimpleNamespace(value=last)
        self._array: np.ndarray = np.zeros(128, dtype=[
            ('index', '<i8'),
            ('time', '<f8'),
            ('close', '<f8'),
        ])
        self._array['index'] = np.arange(128)
        self._array['time'] = np.arange(128)
        self._array['close'] = 100 + np.sin(np.arange(128) / 3)


def test_closed_copy_excludes_mutable_tail_and_detaches_memory(
) -> None:
    '''
    Treating the mutable source tail as finalized leaks provisional
    closes into regime decisions. Make that tail NaN and verify the
    detached prefix remains valid, immutable and independent of a
    later shared-memory write. Absolute first/last bounds survive.

    '''
    source: FakeSource = FakeSource(first=2, last=8)
    source._array['close'][7] = np.nan
    copied: ClosedSourceRows = copy_closed_rows(source)
    assert (copied.first, copied.last, copied.start) == (2, 8, 2)
    assert len(copied.rows) == 5
    assert copied.rows['index'][-1] == 6
    assert not copied.rows.flags.writeable
    original: float = copied.rows['close'][0]
    source._array['close'][2] = 999
    assert copied.rows['close'][0] == original


@pytest.mark.parametrize('first,last', [(0, 8), (2, 6)])
def test_changed_bounds_reject_mixed_copy(
    first: int,
    last: int,
) -> None:
    '''
    A source prepend or truncation between the first bound read and
    array access invalidates the copied range. Publish that change
    from a controlled _array property access; the second bound check
    must reject the candidate before any tracker can consume it.

    '''
    class ChangingSource:
        '''
        Trigger a deterministic source prepend while copying begins.

        '''
        def __init__(self) -> None:
            '''
            Retain a prefix whose origin or last bound will change.

            '''
            self._first: SimpleNamespace = SimpleNamespace(value=2)
            self._last: SimpleNamespace = SimpleNamespace(value=8)
            self.storage: np.ndarray = FakeSource()._array

        @property
        def _array(self) -> np.ndarray:
            '''
            Publish the changed origin after capturing old bounds.

            '''
            self._first.value = first
            self._last.value = last
            return self.storage

    with pytest.raises(SourceChanged, match='changed during copy'):
        copy_closed_rows(ChangingSource())


def test_append_during_copy_preserves_closed_prefix() -> None:
    '''
    A high-frequency producer may append during every large history
    copy. Its old closed prefix remains stable because the captured
    mutable tail was excluded. Accept the detached prefix with its
    original last bound, then consume the newly closed suffix once
    without duplicating observations or starving bootstrap retries.

    '''
    class AppendingSource:
        '''
        Publish two rows after the reader captures initial bounds.

        '''
        def __init__(self) -> None:
            '''
            Retain closed rows before publishing an appended suffix.

            '''
            self._first: SimpleNamespace = SimpleNamespace(value=0)
            self._last: SimpleNamespace = SimpleNamespace(value=6)
            self.storage: np.ndarray = FakeSource()._array

        @property
        def _array(self) -> np.ndarray:
            '''
            Preserve the existing prefix while exposing a new tail.

            '''
            self._last.value = 8
            return self.storage

    source: AppendingSource = AppendingSource()
    copied: ClosedSourceRows = copy_closed_rows(source)
    assert (copied.first, copied.last, copied.start) == (0, 6, 0)
    assert copied.rows['index'].tolist() == [0, 1, 2, 3, 4]
    suffix: ClosedSourceRows = copy_closed_rows(
        source, start=copied.last - 1,
    )
    assert (suffix.first, suffix.last, suffix.start) == (0, 8, 5)
    assert suffix.rows['index'].tolist() == [5, 6]


@pytest.mark.parametrize('first,last', [(-1, 6), (8, 6), (0, 129)])
def test_stable_invalid_bounds_fail_instead_of_spinning(
    first: int,
    last: int,
) -> None:
    '''
    Stable invalid bounds cannot become a usable source through a
    tight retry loop. Treat negative, reversed and out-of-capacity
    ranges as source errors, distinguishing them from bounds that
    changed during the two-read copy validation window.

    '''
    with pytest.raises(ValueError, match='readable range'):
        copy_closed_rows(FakeSource(first, last))


def test_unsigned_coordinate_overflow_is_rejected_before_cast(
) -> None:
    '''
    Converting uint64 coordinates to the core's int64 storage can
    wrap an oversized value negative. Put that value last in a valid
    increasing finalized prefix so ordering checks cannot mask the
    overflow; source validation must reject it before replay begins.

    '''
    source: FakeSource = FakeSource(last=3)
    source._array = source._array.astype([
        ('index', '<u8'),
        ('time', '<f8'),
        ('close', '<f8'),
    ])
    source._array['index'][1] = np.uint64(2 ** 63)
    with pytest.raises(ValueError, match='invalid finalized'):
        copy_closed_rows(source)


@pytest.mark.parametrize('field,value', [
    ('close', float('nan')),
    ('close', float('inf')),
    ('time', float('nan')),
    ('time', 0),
    ('index', -1),
])
def test_bad_finalized_source_is_rejected_before_consumption(
    field: str,
    value: float,
) -> None:
    '''
    Invalid finalized rows must not partially warm a tracker. Corrupt
    a middle row and require validation of the whole detached range
    before it can be handed to the recurrence; mutable-tail exclusion
    is covered separately because unfinished rows may be invalid.

    '''
    source: FakeSource = FakeSource()
    source._array[field][2] = value
    with pytest.raises(ValueError, match='invalid finalized'):
        copy_closed_rows(source)


def test_bootstrap_and_skipped_steps_emit_only_latest_snapshot(
) -> None:
    '''
    A missed sampler wakeup must not skip finalized closes, and a
    bootstrap must not create one HTTP job per historical bar.
    Load five closed rows, mutate the open sixth row, then append
    three rows with one wakeup. Exactly two snapshots must cover the
    complete prefixes; duplicate and foreign messages produce none.

    '''
    published: list[RegimeSnapshot] = []

    async def scenario() -> None:
        '''
        Drive source publication and wakeups in a controlled order.

        '''
        source: FakeSource = FakeSource()
        tracker: MultiTimeframeTrend = MultiTimeframeTrend(
            'test.market', periods=(1, 2),
        )
        send: trio.MemorySendChannel
        receive: trio.MemoryReceiveChannel
        send, receive = trio.open_memory_channel(0)
        first: trio.Event = trio.Event()

        def submit(snapshot: RegimeSnapshot) -> None:
            '''
            Observe publication synchronously at the admission point.

            '''
            published.append(snapshot)
            first.set()

        nursery: trio.Nursery
        async with trio.open_nursery() as nursery:
            nursery.start_soon(
                consume_regime_source,
                source, receive, tracker, submit,
            )
            await first.wait()
            source._array['close'][5] = 500
            await send.send({'period': 1})
            await send.send({'period': 60})
            await send.send({'backfilling': ['other.market', 1]})
            await send.send({'backfilling': ['test.market', 60]})
            await wait_all_tasks_blocked()
            assert len(published) == 1
            source._last.value = 9
            await send.send({'period': 1})
            await send.aclose()
    trio.run(scenario)
    assert len(published) == 2
    assert published[0].sequence == 5
    assert published[1].sequence == 8
    assert published[1].as_of == 8
    assert published[1].revision == 0
    closes: tuple[CloseSample, ...] = (
        published[1].timeframes[0].recent_closes
    )
    assert closes[5].price == 500


@pytest.mark.parametrize(
    'change', ['in_place', 'prepend', 'truncate'],
)
def test_history_changes_revoke_and_replay(change: str) -> None:
    '''
    Appending revised history to warm recurrence state is incorrect,
    including repairs that leave both bounds unchanged. Publish a
    source change and verify immediate revision invalidation precedes
    its replacement snapshot. The rebuilt result must equal replay
    through a fresh tracker over exactly the corrected closed prefix.

    '''
    published: list[RegimeSnapshot] = []
    order: list[tuple[str, int]] = []

    async def scenario() -> None:
        '''
        Compare each revised output with an independent full replay.

        '''
        source: FakeSource = FakeSource(first=2, last=20)
        tracker: MultiTimeframeTrend = MultiTimeframeTrend(
            'test.market', periods=(1, 2),
        )
        send: trio.MemorySendChannel
        receive: trio.MemoryReceiveChannel
        send, receive = trio.open_memory_channel(0)
        first: trio.Event = trio.Event()

        def submit(snapshot: RegimeSnapshot) -> None:
            '''
            Preserve publication order and source generation.

            '''
            published.append(snapshot)
            order.append(('snapshot', snapshot.revision))
            first.set()

        def invalidate(revision: int) -> None:
            '''
            Record when old pending judgments become inadmissible.

            '''
            order.append(('invalidate', revision))

        async def consume() -> None:
            '''
            Attach both source callbacks to one consumer lifetime.

            '''
            await consume_regime_source(
                source, receive, tracker, submit,
                invalidate=invalidate,
            )

        nursery: trio.Nursery
        async with trio.open_nursery() as nursery:
            nursery.start_soon(consume)
            await first.wait()
            if change == 'in_place':
                source._array['close'][4] = 200
                await send.send({
                    'backfilling': ['test.market', 1],
                })
            else:
                if change == 'prepend':
                    source._first.value = 0
                else:
                    source._last.value = 10
                await send.send({'period': 1})
            await send.aclose()
        reference: MultiTimeframeTrend = MultiTimeframeTrend(
            'test.market', periods=(1, 2), revision=1,
        )
        row: np.void
        for row in copy_closed_rows(source).rows:
            expected: RegimeSnapshot = reference.push(
                float(row['time']), float(row['close']),
                index=int(row['index']),
            )
        assert published[-1] == expected

    trio.run(scenario)
    assert order == [
        ('snapshot', 0), ('invalidate', 1), ('snapshot', 1),
    ]


def test_revision_during_bootstrap_suppresses_old_publication(
) -> None:
    '''
    Losing a backfill during bootstrap publishes a mixed-generation
    classification that ordinary bound checks cannot detect.
    ``SignallingTracker.push()`` sets ``entered`` after one row;
    ``consume_regime_source()`` yields at every replay row so the
    ``scenario()`` task can repair a close and send its notice.
    ``_Notices.collect()`` must invalidate that revision before the
    replay publishes. One replacement snapshot, its corrected close
    and full sequence prove old evidence never reaches ``published``.

    '''
    published: list[RegimeSnapshot] = []
    invalidated: list[int] = []

    async def scenario() -> None:
        '''
        Synchronize repair with the first chunk's cooperative yield.

        '''
        entered: trio.Event = trio.Event()

        class SignallingTracker(MultiTimeframeTrend):
            '''
            Expose the start of bootstrap without adding I/O to push.

            '''
            def push(
                self,
                *args: Any,
                **kwargs: Any,
            ) -> RegimeSnapshot:
                '''
                Signal after the first historical row is consumed.

                '''
                snapshot: RegimeSnapshot = super().push(
                    *args, **kwargs,
                )
                entered.set()
                return snapshot

        source: FakeSource = FakeSource(last=100)
        tracker: SignallingTracker = SignallingTracker(
            'test.market', periods=(1,), history_limit=128,
        )
        send: trio.MemorySendChannel
        receive: trio.MemoryReceiveChannel
        send, receive = trio.open_memory_channel(0)

        async def consume() -> None:
            '''
            Yield each row so the controlled repair reaches replay.

            '''
            await consume_regime_source(
                source, receive, tracker, published.append,
                invalidate=invalidated.append,
                replay_chunk=1,
            )

        nursery: trio.Nursery
        async with trio.open_nursery() as nursery:
            nursery.start_soon(consume)
            await entered.wait()
            source._array['close'][0] = 999
            await send.send({'backfilling': ['test.market', 1]})
            await send.aclose()

    trio.run(scenario)
    assert invalidated == [1]
    assert len(published) == 1
    assert published[0].revision == 1
    assert published[0].sequence == 99
    assert published[0].timeframes[0].recent_closes[0].price == 999


@pytest.mark.parametrize('quote_lag', [False, True])
@pytest.mark.parametrize('source_age', [None, 0, 60, -10])
def test_native_wrapper_subscribes_filters_age_and_drains_quotes(
    monkeypatch: pytest.MonkeyPatch,
    quote_lag: bool,
    source_age: float|None,
) -> None:
    '''
    Subscribing after history capture loses in-place revisions in the
    startup gap. ``GuardedSource._array`` asserts the subscription is
    active before every read. ``sampler_context()`` then replaces
    ``Flume.stream`` with ``OtherConsumer``; any later lookup of that
    shared field fails instead of stealing another consumer's stream.

    The independent cases cover these admission contracts:
    - ``LaggingQuotes`` raises ``Lagged`` once. It sets ``recovered``
      only when the quote drain resumes; SHM processing must survive.
    - Old/future source clocks must withhold model evidence even
      though submission occurs locally now; no age limit permits
      offline replay. ``published`` distinguishes those outcomes.
    - ``feed_context()`` rejects a requested quote throttle: feeds
      cached solely by market must retain their shared quote policy.

    ``finish()`` waits on ``read_source`` and, when needed,
    ``recovered``, then closes sampling after tasks block. Context
    exit order proves the native wrapper releases its subscription
    before the feed, without any broker, actor service or API key.

    '''
    order: list[str] = []
    subscribed: bool = False
    published: list[RegimeSnapshot] = []
    read_source: trio.Event
    flume_ref: SimpleNamespace|None = None

    class GuardedSource:
        '''
        Require the native sampler subscription before source access.

        '''
        def __init__(self) -> None:
            '''
            Seed a readable range whose array access is guarded.

            '''
            self._first: SimpleNamespace = SimpleNamespace(value=0)
            self._last: SimpleNamespace = SimpleNamespace(value=6)
            self.storage: np.ndarray = FakeSource()._array

        @property
        def _array(self) -> np.ndarray:
            '''
            Detect a source read before entering the subscription.

            '''
            assert subscribed
            order.append('read_source')
            read_source.set()
            return self.storage

    async def scenario() -> None:
        '''
        Close the synthetic sampler to end the actual wrapper.

        '''
        nonlocal read_source
        read_source = trio.Event()
        recovered: trio.Event = trio.Event()
        send: trio.MemorySendChannel
        receive: trio.MemoryReceiveChannel
        send, receive = trio.open_memory_channel(0)
        quote_send: trio.MemorySendChannel
        quote_receive: trio.MemoryReceiveChannel
        quote_send, quote_receive = trio.open_memory_channel(0)

        class LaggingQuotes:
            '''
            Lose one discarded quote before resuming a healthy wait.

            '''
            def __aiter__(self) -> 'LaggingQuotes':
                '''
                Return the resumable stream after a Lagged exception.

                '''
                return self

            async def __anext__(self) -> dict:
                '''
                Force one overrun, then prove the drain resumed.

                '''
                if 'quote_lagged' not in order:
                    order.append('quote_lagged')
                    raise Lagged('synthetic discarded quote')
                recovered.set()
                await trio.sleep_forever()

        @asynccontextmanager
        async def feed_context(
            **kwargs: object,
        ) -> AsyncIterator[SimpleNamespace]:
            '''
            Expose a source Flume with a cancellable quote stream.

            '''
            nonlocal flume_ref
            assert 'tick_throttle' not in kwargs
            order.append('open_feed')
            flume_ref = SimpleNamespace(
                rt_sample_period_s=1,
                rt_shm=GuardedSource(),
                stream=(
                    LaggingQuotes() if quote_lag else quote_receive
                ),
            )
            yield SimpleNamespace(flumes={
                'test.market': flume_ref,
            })
            await quote_send.aclose()
            order.append('close_feed')

        @asynccontextmanager
        async def sampler_context(
            **kwargs: object,
        ) -> AsyncIterator[trio.MemoryReceiveChannel]:
            '''
            Mark exactly when the sampler subscription begins.

            '''
            nonlocal subscribed
            subscribed = True
            order.append('subscribe')

            class OtherConsumer:
                '''
                Model a shared Flume's receiver ownership change.

                '''
                def __aiter__(self) -> Never:
                    '''
                    Detect consuming a later caller's broadcast view.

                    '''
                    pytest.fail('drain stole another quote receiver')

            assert flume_ref is not None
            flume_ref.stream = OtherConsumer()
            try:
                yield receive
            finally:
                subscribed = False
                order.append('unsubscribe')

        monkeypatch.setattr(
            'piker.data.feed.maybe_open_feed', feed_context,
        )
        monkeypatch.setattr(
            'piker.data._sampling.open_sample_stream',
            sampler_context,
        )
        if source_age is not None:
            monkeypatch.setattr(
                'piker.fsp._regime_feed.time.time',
                lambda: 5 + source_age,
            )

        async def finish() -> None:
            '''
            Close the sampler after processing and draining quotes.

            '''
            await read_source.wait()
            if quote_lag:
                await recovered.wait()
            await wait_all_tasks_blocked()
            await send.aclose()

        nursery: trio.Nursery
        async with trio.open_nursery() as nursery:
            nursery.start_soon(finish)
            await run_regime_feed(
                'test.market', published.append,
                max_source_age=None if source_age is None else 5,
            )

    trio.run(scenario)
    assert len(published) == (source_age in (None, 0))
    assert order.index('subscribe') < order.index('read_source')
    assert order[-2:] == ['unsubscribe', 'close_feed']


def test_service_composes_feed_and_classifier_in_one_lifetime(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    '''
    A prepared provider adapter needs concrete feed-to-classifier
    wiring and shared cancellation ownership. Replace the broker
    source with a synthetic callback producer but retain the actual
    HTTP client, worker and service nursery. One corrected snapshot
    must produce one validated judgment; normal source exit must
    stop the classifier while leaving the caller's HTTP session open.

    '''
    from piker.fsp._jev import (
        DEFAULT_MODEL,
        JevClassification,
        JevClient,
        RegimeClassifier,
    )

    published: list[JevClassification] = []

    async def scenario() -> None:
        '''
        End the source only after its judgment is published.

        '''
        delivered: trio.Event = trio.Event()

        async def source(
            fqme: str,
            submit: Callable[[RegimeSnapshot], object],
            *,
            invalidate: Callable[[int], object],
            **kwargs: object,
        ) -> None:
            '''
            Replay one revision through the bound service callbacks.

            '''
            assert fqme == 'test.market'
            assert kwargs['max_source_age'] == 5
            invalidate(1)
            tracker: MultiTimeframeTrend = MultiTimeframeTrend(
                fqme, periods=(1,), revision=1,
            )
            second: int
            for second in range(8):
                snapshot: RegimeSnapshot = tracker.push(
                    second, second,
                )
            submit(snapshot)
            await delivered.wait()

        def handler(request: httpx.Request) -> httpx.Response:
            '''
            Return a documented Choice response without real HTTP.

            '''
            request_body: dict = json.loads(request.content)
            answers: dict[str, dict] = {}
            name: str
            question: dict
            for name, question in request_body['questions'].items():
                probabilities: dict[str, float] = {}
                label: str
                for label in question['criteria']:
                    probabilities[label] = float(label == 'uptrend')
                answers[name] = {
                    'type': 'choice',
                    'choice': 'uptrend',
                    'probabilities': probabilities,
                    'confidence': 1,
                }
            return httpx.Response(200, json={
                'model': DEFAULT_MODEL,
                'answers': answers,
                'usage': {
                    'input_tokens': 100,
                    'output_tokens': 20,
                },
            })

        async def publish(result: JevClassification) -> None:
            '''
            Release source shutdown after receiving a typed judgment.

            '''
            published.append(result)
            delivered.set()

        monkeypatch.setattr(
            'piker.fsp._regime_feed.run_regime_feed', source,
        )
        http: httpx.AsyncClient
        async with httpx.AsyncClient(
            transport=httpx.MockTransport(handler),
        ) as http:
            client: JevClient = JevClient(http, 'offline-test-key')
            worker: RegimeClassifier = await serve_regimes(
                'test.market', client, publish,
            )
            assert not http.is_closed
            assert worker.requests == 1
            assert worker.published == 1
        assert http.is_closed

    trio.run(scenario)
    assert len(published) == 1
    assert published[0].market == 'test.market'
    assert published[0].revision == 1