piker/tests/providers/test_ib_schedule.py

# piker: trading gear for hackers
# Copyright (C) 2018-present  Tyler Goodlet (in stewardship of pikers)

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

'''
Historical schedule normalization and request-boundary regressions.

'''
import asyncio
from datetime import UTC, datetime
from types import SimpleNamespace
from zoneinfo import ZoneInfo

import pytest

from piker.brokers.ib.venues import normalize_schedule
from piker.brokers.ib.api import Client


def response(zone='America/New_York'):
    '''
    Describe two sessions separated by a maintenance closure.

    '''
    return SimpleNamespace(
        timeZone=zone,
        startDateTime='20260923-00:00:00',
        endDateTime='20260925-00:00:00',
        sessions=[SimpleNamespace(
            startDateTime='20260923-18:00:00',
            endDateTime='20260924-17:00:00',
        )],
    )


def test_schedule_uses_dated_timezone_and_coverage():
    '''
    Historical allocation must use dated sessions, not today's hours.

    A New York September session must convert using EDT and preserve
    the response's wider coverage separately from its open interval.

    '''
    out = normalize_schedule(response(), False)
    expected = datetime(2026, 9, 23, 22, tzinfo=UTC).timestamp()
    assert out['sessions'][0][0] == expected
    assert out['start'] < expected < out['end']
    assert out['use_rth'] is False


@pytest.mark.parametrize('stamp', [
    '20261101-01:30:00', '20260308-02:30:00',
])
def test_schedule_rejects_ambiguous_or_nonexistent_time(stamp):
    '''
    A guessed DST fold can shift reservations by an hour.

    Reject both the repeated autumn hour and nonexistent spring
    hour instead of turning unknown coverage into a false closure.

    '''
    schedule = response()
    schedule.startDateTime = stamp
    with pytest.raises(ValueError, match='Ambiguous'):
        normalize_schedule(schedule, False)


def test_schedule_request_uses_extended_hours_and_utc():
    '''
    Schedule requests must match the bars endpoint's extended hours.

    Exercise the real Client method with a fake IB transport. Assert
    exact UTC request bounds and day rounding, and ensure response
    coverage is not replaced by the caller's requested interval.

    '''
    calls = []
    async def find_contracts(fqme):
        assert fqme == 'mnq.cme.20261218'
        return ['contract']
    async def schedule(contract, **kwargs):
        calls.append((contract, kwargs))
        return response()
    client = SimpleNamespace(
        find_contracts=find_contracts,
        ib=SimpleNamespace(reqHistoricalScheduleAsync=schedule),
    )
    out = asyncio.run(Client.history_schedule(
        client, 'mnq.cme.20261218',
        datetime(2026, 9, 23, 4, tzinfo=UTC),
        datetime(2026, 9, 24, 1, tzinfo=UTC),
    ))
    assert calls == [('contract', {
        'numDays': 1, 'endDateTime': '20260924-01:00:00',
        'useRTH': False,
    })]
    assert out['timezone'] == 'America/New_York'


def test_schedule_timeout_is_not_an_empty_session_list():
    '''
    A failed schedule request must not classify the range as closed.

    Block the simulated provider and apply an immediate timeout;
    the adapter must propagate failure instead of returning coverage.

    '''
    async def find_contracts(fqme):
        return ['contract']
    async def schedule(*args, **kwargs):
        await asyncio.sleep(100)
    client = SimpleNamespace(
        find_contracts=find_contracts,
        ib=SimpleNamespace(reqHistoricalScheduleAsync=schedule),
    )
    with pytest.raises(TimeoutError):
        asyncio.run(Client.history_schedule(
            client, 'mnq.cme',
            datetime(2026, 9, 23, tzinfo=UTC),
            datetime(2026, 9, 24, tzinfo=UTC),
            timeout_s=0,
        ))


def test_schedule_rejects_reversed_session():
    '''
    A malformed session must not silently become a venue closure.

    Reject the response instead of dropping its reversed interval
    and reserving slots using the remaining incomplete session list.

    '''
    schedule = response()
    schedule.sessions[0].endDateTime = '20260923-17:00:00'
    with pytest.raises(ValueError, match='Invalid schedule session'):
        normalize_schedule(schedule, False)


def test_schedule_duration_counts_elapsed_time_across_dst():
    '''
    The autumn clock change makes this local day last 25 hours.

    Request two elapsed days so duration rounding cannot leave an
    hour uncovered by subtracting same-zone wall-clock datetimes.

    '''
    async def find_contracts(fqme):
        return ['contract']

    async def schedule(contract, **kwargs):
        assert kwargs['numDays'] == 2
        return response()

    client = SimpleNamespace(
        find_contracts=find_contracts,
        ib=SimpleNamespace(reqHistoricalScheduleAsync=schedule),
    )
    zone = ZoneInfo('America/New_York')
    asyncio.run(Client.history_schedule(
        client, 'mnq.cme',
        datetime(2026, 11, 1, tzinfo=zone),
        datetime(2026, 11, 2, tzinfo=zone),
    ))


def mnq_schedule(wider: bool = False) -> SimpleNamespace:
    '''
    Reproduce the dated US/Central sessions captured from live IB.

    '''
    starts: list[str] = ['20260923-17:00:00']
    ends: list[str] = ['20260924-16:00:00']
    if wider:
        starts.insert(0, '20260922-17:00:00')
        ends.insert(0, '20260923-16:00:00')
    sessions: list[SimpleNamespace] = []
    start: str
    end: str
    for start, end in zip(starts, ends):
        sessions.append(SimpleNamespace(
            startDateTime=start,
            endDateTime=end,
        ))
    return SimpleNamespace(
        timeZone='US/Central',
        startDateTime=starts[0],
        endDateTime=ends[-1],
        sessions=sessions,
    )


def test_schedule_retry_reserves_mnq_gap_outside_closure():
    '''
    One elapsed day returned only MNQ's newest trading session.

    Its coverage began after the left gap margin, so the allocator
    could not reserve any slots. Replay both captured IB responses
    through Client.history_schedule() and reserve_history_gaps().
    Assert that widening keeps the end fixed, reserves both trading
    portions, and leaves the one-hour maintenance closure compressed.
    This uses fake transport and arrays, with no live SHM or storage.

    '''
    import numpy as np
    from piker.tsp._reserve import reserve_history_gaps

    calls: list[dict] = []

    async def find_contracts(fqme: str) -> list[str]:
        return ['contract']

    async def schedule(contract: str, **kwargs) -> SimpleNamespace:
        calls.append(kwargs)
        return mnq_schedule(wider=len(calls) > 1)

    client = SimpleNamespace(
        find_contracts=find_contracts,
        ib=SimpleNamespace(reqHistoricalScheduleAsync=schedule),
    )

    async def query(left: float, right: float) -> dict:
        return await Client.history_schedule(
            client, 'mnq.cme.20261218',
            datetime.fromtimestamp(left, UTC),
            datetime.fromtimestamp(right, UTC),
        )

    frame: np.ndarray = np.zeros(2, dtype=[
        ('time', 'f8'),
        ('open', 'f8'),
        ('high', 'f8'),
        ('low', 'f8'),
        ('close', 'f8'),
    ])
    frame['time'] = [1790193489, 1790248914]
    frame['close'] = [30755.75, 30478.5]
    result: np.ndarray
    ranges: list[tuple[float, float]]
    result, ranges = asyncio.run(reserve_history_gaps(
        frame, 1, query, max_rows=51826,
    ))
    assert calls == [
        {
            'numDays': days,
            'endDateTime': '20260924-11:21:54',
            'useRTH': False,
        }
        for days in (1, 2)
    ]
    assert len(result) == 51826
    assert ranges == [
        (1790193490, 1790197200),
        (1790200800, 1790248914),
    ]
    assert not np.any(
        (result['time'] >= 1790197200)
        & (result['time'] < 1790200800)
    )
    np.testing.assert_array_equal(result[[0, -1]], frame)


@pytest.mark.parametrize('missing_right', [False, True])
def test_schedule_retry_preserves_unknown_coverage(
    missing_right: bool,
):
    '''
    Widening must not turn a partial reply into claimed coverage.

    Return the same incomplete schedule on every fake request.
    Missing left coverage exhausts the bounded lookback; missing
    right coverage stops immediately because an earlier start cannot
    repair it. The unchanged response retains the unknown interval
    so downstream reservation still rejects it.

    '''
    calls: list[int] = []

    async def find_contracts(fqme: str) -> list[str]:
        return ['contract']

    async def schedule(contract: str, **kwargs) -> SimpleNamespace:
        calls.append(kwargs['numDays'])
        return mnq_schedule()

    client = SimpleNamespace(
        find_contracts=find_contracts,
        ib=SimpleNamespace(reqHistoricalScheduleAsync=schedule),
    )
    end: datetime = datetime(2026, 9, 24, 11, tzinfo=UTC)
    if missing_right:
        end = datetime(2026, 9, 24, 23, tzinfo=UTC)
    out: dict = asyncio.run(Client.history_schedule(
        client, 'mnq.cme.20261218',
        datetime(2026, 9, 23, 19, tzinfo=UTC), end,
    ))
    assert out == normalize_schedule(mnq_schedule(), False)
    assert calls == ([2] if missing_right else [1, 2, 4, 8, 16, 31])


def test_schedule_retry_shares_timeout(
    monkeypatch: pytest.MonkeyPatch,
):
    '''
    A retry must not renew the caller's total schedule wait budget.

    Record timeout scopes while the first fake reply lacks coverage
    and the second request waits forever. The shared deadline must
    cancel that second request and propagate TimeoutError, with no
    partial result substituted for failure. No sleep ordering is
    required: only the blocked second request yields to the deadline.

    '''
    timeouts: list[float] = []
    calls: int = 0
    cancelled: bool = False
    original_timeout = asyncio.timeout

    def timeout(delay: float):
        timeouts.append(delay)
        return original_timeout(delay)

    monkeypatch.setattr(asyncio, 'timeout', timeout)

    async def find_contracts(fqme: str) -> list[str]:
        return ['contract']

    async def schedule(contract: str, **kwargs) -> SimpleNamespace:
        nonlocal calls, cancelled
        calls += 1
        if calls == 1:
            return mnq_schedule()
        try:
            await asyncio.Future()
        except asyncio.CancelledError:
            cancelled = True
            raise

    client = SimpleNamespace(
        find_contracts=find_contracts,
        ib=SimpleNamespace(reqHistoricalScheduleAsync=schedule),
    )
    with pytest.raises(TimeoutError):
        asyncio.run(Client.history_schedule(
            client, 'mnq.cme.20261218',
            datetime(2026, 9, 23, 19, tzinfo=UTC),
            datetime(2026, 9, 24, 11, tzinfo=UTC),
            timeout_s=0.01,
        ))
    assert timeouts == [0.01]
    assert calls == 2
    assert cancelled