tests/test_data_layer.py

import math
import unittest
from datetime import timedelta
from unittest import mock
import numpy as np
import pandas as pd
from config import YF_NATIVE_INTERVALS
from data.historical import OHLCV_COLUMNS, DataFetchError, get_historical, get_historical_with_features, resample_ohlc
from features.feature_pipeline import FEATURE_COLUMNS

def _window_candles(start, end, freq: str) -> pd.DataFrame:
    idx = pd.date_range(start=start, end=end, freq=freq, inclusive='left')
    if len(idx) == 0:
        return pd.DataFrame()
    n = len(idx)
    base = 100.0 + np.arange(n, dtype=float) * 0.1
    return pd.DataFrame({'Open': base, 'High': base + 0.5, 'Low': base - 0.5, 'Close': base + 0.2, 'Volume': np.arange(n, dtype=float) + 1.0, 'Dividends': np.zeros(n)}, index=idx)

class _RecordingFetcher:

    def __init__(self, freq: str='1h', empty_calls=()):
        self.freq = freq
        self.empty_calls = set(empty_calls)
        self.calls = []
        self.frames = []

    def __call__(self, symbol, interval, start, end):
        self.calls.append({'symbol': symbol, 'interval': interval, 'start': start, 'end': end})
        if len(self.calls) in self.empty_calls:
            return pd.DataFrame()
        frame = _window_candles(start, end, self.freq)
        self.frames.append(frame)
        return frame

    @property
    def spans(self):
        return [c['end'] - c['start'] for c in self.calls]

    @property
    def total_span(self):
        return max((c['end'] for c in self.calls)) - min((c['start'] for c in self.calls))

    def raw(self) -> pd.DataFrame:
        return pd.concat(self.frames).sort_index()

class TestLookbackIsHonored(unittest.TestCase):

    def test_daily_lookback_requests_exactly_that_window(self):
        fetcher = _RecordingFetcher(freq='1D')
        with mock.patch('data.historical._fetch_yf', fetcher):
            df = get_historical('BTC-USD', '1d', lookback_days=30)
        self.assertEqual(len(fetcher.calls), 1)
        self.assertEqual(fetcher.calls[0]['interval'], '1d')
        self.assertEqual(fetcher.spans[0], timedelta(days=30))
        self.assertLessEqual(len(df), 31)
        self.assertGreater(len(df), 0)

    def test_returned_frame_is_ordered_and_ohlcv_only(self):
        fetcher = _RecordingFetcher(freq='1D')
        with mock.patch('data.historical._fetch_yf', fetcher):
            df = get_historical('BTC-USD', '1d', lookback_days=10)
        self.assertEqual(list(df.columns), OHLCV_COLUMNS)
        self.assertNotIn('Dividends', df.columns)
        self.assertTrue(df.index.is_monotonic_increasing)
        self.assertFalse(df.index.duplicated().any())

    def test_no_lookback_falls_back_to_the_interval_limit(self):
        fetcher = _RecordingFetcher(freq='1h')
        with mock.patch('data.historical._fetch_yf', fetcher):
            get_historical('BTC-USD', '60m')
        self.assertEqual(fetcher.total_span, timedelta(days=YF_NATIVE_INTERVALS['60m']['max_days']))

class TestChunkedFetching(unittest.TestCase):

    def setUp(self):
        self.limits = YF_NATIVE_INTERVALS['1m']

    def test_window_wider_than_the_request_limit_is_split(self):
        requested = 30
        fetcher = _RecordingFetcher(freq='1h')
        with mock.patch('data.historical._fetch_yf', fetcher):
            df = get_historical('BTC-USD', '1m', lookback_days=requested)
        request_days = self.limits['max_request_days']
        self.assertEqual(len(fetcher.calls), math.ceil(requested / request_days))
        for span in fetcher.spans:
            self.assertLessEqual(span, timedelta(days=request_days))
        self.assertEqual(sum(fetcher.spans, timedelta()), timedelta(days=requested))
        self.assertEqual(fetcher.total_span, timedelta(days=requested))
        self.assertGreater(len(df), 0)

    def test_chunks_are_contiguous_with_no_gap_or_overlap(self):
        fetcher = _RecordingFetcher(freq='1h')
        with mock.patch('data.historical._fetch_yf', fetcher):
            get_historical('BTC-USD', '1m', lookback_days=20)
        for newer, older in zip(fetcher.calls, fetcher.calls[1:]):
            self.assertEqual(newer['start'], older['end'])

    def test_lookback_beyond_the_source_limit_is_clamped(self):
        fetcher = _RecordingFetcher(freq='1h')
        with mock.patch('data.historical._fetch_yf', fetcher):
            get_historical('BTC-USD', '1m', lookback_days=90)
        max_days = self.limits['max_days']
        self.assertEqual(fetcher.total_span, timedelta(days=max_days))
        self.assertEqual(len(fetcher.calls), math.ceil(max_days / self.limits['max_request_days']))

    def test_one_empty_chunk_does_not_abort_the_fetch(self):
        fetcher = _RecordingFetcher(freq='1h', empty_calls=(2,))
        with mock.patch('data.historical._fetch_yf', fetcher):
            df = get_historical('BTC-USD', '1m', lookback_days=30)
        self.assertEqual(len(fetcher.calls), 5)
        self.assertEqual(len(df), len(fetcher.raw()))
        self.assertGreater(len(df), 0)

    def test_all_chunks_empty_reports_the_source_limit(self):
        fetcher = _RecordingFetcher(freq='1h', empty_calls=range(1, 6))
        with mock.patch('data.historical._fetch_yf', fetcher):
            with self.assertRaises(DataFetchError) as ctx:
                get_historical('NOPE-USD', '1m', lookback_days=30)
        message = str(ctx.exception)
        self.assertIn('NOPE-USD', message)
        self.assertIn('1m', message)
        self.assertIn(str(self.limits['max_days']), message)

class TestTimeframeValidation(unittest.TestCase):

    def test_unknown_timeframe_raises_before_any_request(self):
        fetcher = _RecordingFetcher()
        with mock.patch('data.historical._fetch_yf', fetcher):
            with self.assertRaises(DataFetchError):
                get_historical('BTC-USD', '3y', lookback_days=5)
        self.assertEqual(fetcher.calls, [])

    def test_blank_symbol_raises_before_any_request(self):
        fetcher = _RecordingFetcher()
        with mock.patch('data.historical._fetch_yf', fetcher):
            with self.assertRaises(DataFetchError):
                get_historical('   ', '1d', lookback_days=5)
        self.assertEqual(fetcher.calls, [])

class TestResampling(unittest.TestCase):

    def test_custom_timeframe_resamples_from_its_source_interval(self):
        fetcher = _RecordingFetcher(freq='1h')
        with mock.patch('data.historical._fetch_yf', fetcher):
            df = get_historical('BTC-USD', '4h', lookback_days=5)
        self.assertEqual([c['interval'] for c in fetcher.calls], ['60m'])
        spacing = df.index.to_series().diff().dropna().unique()
        self.assertEqual(list(spacing), [pd.Timedelta(hours=4)])
        raw = fetcher.raw()
        bar_time = df.index[len(df) // 2]
        bucket = raw[(raw.index >= bar_time) & (raw.index < bar_time + pd.Timedelta(hours=4))]
        self.assertEqual(len(bucket), 4)
        self.assertAlmostEqual(df.loc[bar_time, 'Open'], float(bucket['Open'].iloc[0]))
        self.assertAlmostEqual(df.loc[bar_time, 'High'], float(bucket['High'].max()))
        self.assertAlmostEqual(df.loc[bar_time, 'Low'], float(bucket['Low'].min()))
        self.assertAlmostEqual(df.loc[bar_time, 'Close'], float(bucket['Close'].iloc[-1]))
        self.assertAlmostEqual(df.loc[bar_time, 'Volume'], float(bucket['Volume'].sum()))

    def test_resample_ohlc_aggregates_each_field(self):
        idx = pd.date_range('2024-01-01', periods=4, freq='1h')
        df = pd.DataFrame({'Open': [10.0, 11.0, 12.0, 13.0], 'High': [15.0, 11.5, 20.0, 13.5], 'Low': [9.0, 8.0, 11.0, 12.0], 'Close': [11.0, 12.0, 13.0, 14.0], 'Volume': [1.0, 2.0, 3.0, 4.0]}, index=idx)
        out = resample_ohlc(df, '4h')
        self.assertEqual(len(out), 1)
        row = out.iloc[0]
        self.assertEqual(row['Open'], 10.0)
        self.assertEqual(row['High'], 20.0)
        self.assertEqual(row['Low'], 8.0)
        self.assertEqual(row['Close'], 14.0)
        self.assertEqual(row['Volume'], 10.0)

class TestGetHistoricalWithFeatures(unittest.TestCase):

    def test_features_off_returns_none_and_requests_no_padding(self):
        fetcher = _RecordingFetcher(freq='1h')
        with mock.patch('data.historical._fetch_yf', fetcher):
            df, features = get_historical_with_features('BTC-USD', '1h', lookback_days=3, use_features=False)
        self.assertIsNone(features)
        self.assertEqual(fetcher.total_span, timedelta(days=3))
        self.assertGreater(len(df), 0)

    def test_features_on_pads_history_and_returns_aligned_clean_frames(self):
        fetcher = _RecordingFetcher(freq='1h')
        with mock.patch('data.historical._fetch_yf', fetcher):
            df, features = get_historical_with_features('BTC-USD', '1h', lookback_days=3, use_features=True)
        self.assertGreater(fetcher.total_span, timedelta(days=3))
        self.assertEqual(list(features.columns), FEATURE_COLUMNS)
        self.assertEqual(len(features), len(df))
        self.assertFalse(features.isna().any().any())
        self.assertLess(len(df), len(fetcher.raw()))
if __name__ == '__main__':
    unittest.main()