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