tests/test_helpers.py

import unittest
import pandas as pd
from config import CUSTOM_TIMEFRAMES, YF_NATIVE_INTERVALS
from utils.helpers import INDICATOR_WARMUP_CANDLES, MINUTES_PER_CANDLE, PREDICTION_HISTORY_CANDLES, drop_forming_candle, feature_warmup_days, future_index, history_days_for_candles, infer_step

def _frame(index) -> pd.DataFrame:
    return pd.DataFrame({'Close': range(len(index))}, index=index)

class TestHistoryDaysForCandles(unittest.TestCase):

    def test_window_covers_the_requested_candle_count(self):
        for timeframe, minutes in MINUTES_PER_CANDLE.items():
            days = history_days_for_candles(timeframe, 200)
            self.assertGreaterEqual(days * 24 * 60, 200 * minutes, timeframe)

    def test_known_sizes(self):
        self.assertEqual(history_days_for_candles('1d', 200), 200)
        self.assertEqual(history_days_for_candles('4h', 200), 34)
        self.assertEqual(history_days_for_candles('5m', 200), 1)

    def test_minimum_days_acts_as_a_floor_only(self):
        self.assertEqual(history_days_for_candles('1m', 10, minimum_days=2), 2)
        self.assertEqual(history_days_for_candles('1d', 10, minimum_days=2), 10)

    def test_non_positive_candle_counts_still_request_a_real_window(self):
        self.assertEqual(history_days_for_candles('1h', 0), 1)
        self.assertEqual(history_days_for_candles('1h', -5), 1)

    def test_unknown_timeframe_uses_the_hourly_fallback(self):
        self.assertEqual(history_days_for_candles('7m', 200), history_days_for_candles('1h', 200))

class TestFeatureWarmupDays(unittest.TestCase):

    def test_padding_covers_the_indicator_warmup(self):
        for timeframe, minutes in MINUTES_PER_CANDLE.items():
            days = feature_warmup_days(timeframe, minimum_days=0)
            self.assertGreaterEqual(days * 24 * 60, INDICATOR_WARMUP_CANDLES * minutes, timeframe)

    def test_default_floor_is_two_days(self):
        self.assertEqual(feature_warmup_days('1m'), 2)

    def test_floor_can_be_dropped_for_padding_math(self):
        self.assertEqual(feature_warmup_days('1m', minimum_days=0), 1)
        self.assertEqual(feature_warmup_days('1d', minimum_days=0), 30)

class TestLiveSizingFitsSourceLimits(unittest.TestCase):

    def test_prediction_history_is_servable_for_every_timeframe(self):
        for timeframe in MINUTES_PER_CANDLE:
            source = CUSTOM_TIMEFRAMES.get(timeframe, {}).get('source', timeframe)
            limit = YF_NATIVE_INTERVALS.get(source, {}).get('max_days')
            if limit is None:
                continue
            wanted = history_days_for_candles(timeframe, PREDICTION_HISTORY_CANDLES)
            self.assertLessEqual(wanted, limit, f"{timeframe} wants {wanted}d but '{source}' serves {limit}d")

class TestInferStep(unittest.TestCase):

    def test_uses_the_median_spacing(self):
        idx = pd.date_range('2024-01-01', periods=5, freq='15min')
        self.assertEqual(infer_step(idx), pd.Timedelta(minutes=15))

    def test_market_closure_gap_does_not_dominate(self):
        idx = pd.DatetimeIndex(['2024-01-05 20:00', '2024-01-05 21:00', '2024-01-05 22:00', '2024-01-08 20:00', '2024-01-08 21:00'])
        self.assertEqual(infer_step(idx), pd.Timedelta(hours=1))

    def test_single_candle_falls_back_to_one_minute(self):
        idx = pd.date_range('2024-01-01', periods=1, freq='5min')
        self.assertEqual(infer_step(idx), pd.Timedelta(minutes=1))

class TestFutureIndex(unittest.TestCase):

    def test_projects_timestamps_after_the_last_candle(self):
        last = pd.Timestamp('2024-01-01 00:00')
        self.assertEqual(future_index(last, pd.Timedelta(minutes=5), 3), [pd.Timestamp('2024-01-01 00:05'), pd.Timestamp('2024-01-01 00:10'), pd.Timestamp('2024-01-01 00:15')])

    def test_zero_horizon_projects_nothing(self):
        self.assertEqual(future_index(pd.Timestamp('2024-01-01'), pd.Timedelta(minutes=1), 0), [])

class TestDropFormingCandle(unittest.TestCase):

    def test_drops_the_still_forming_candle_and_its_feature_row(self):
        now = pd.Timestamp.now(tz='UTC').floor('min')
        idx = pd.DatetimeIndex([now - pd.Timedelta(minutes=10), now - pd.Timedelta(minutes=5), now])
        features = pd.DataFrame({'f': [1.0, 2.0, 3.0]})
        df, trimmed_features = drop_forming_candle(_frame(idx), features)
        self.assertEqual(len(df), 2)
        self.assertEqual(df.index[-1], idx[-2])
        self.assertEqual(len(trimmed_features), 2)
        self.assertEqual(list(trimmed_features.index), [0, 1])

    def test_keeps_a_closed_final_candle(self):
        now = pd.Timestamp.now(tz='UTC').floor('min')
        idx = pd.DatetimeIndex([now - pd.Timedelta(minutes=20), now - pd.Timedelta(minutes=15), now - pd.Timedelta(minutes=10)])
        features = pd.DataFrame({'f': [1.0, 2.0, 3.0]})
        df, kept_features = drop_forming_candle(_frame(idx), features)
        self.assertEqual(len(df), 3)
        self.assertEqual(len(kept_features), 3)

    def test_handles_a_tz_naive_index(self):
        now = pd.Timestamp.now(tz='UTC').tz_localize(None).floor('min')
        idx = pd.DatetimeIndex([now - pd.Timedelta(minutes=10), now - pd.Timedelta(minutes=5), now])
        df, features = drop_forming_candle(_frame(idx), None)
        self.assertEqual(len(df), 2)
        self.assertIsNone(features)

    def test_empty_frame_passes_through(self):
        empty = pd.DataFrame({'Close': []}, index=pd.DatetimeIndex([]))
        df, features = drop_forming_candle(empty, None)
        self.assertEqual(len(df), 0)
        self.assertIsNone(features)
if __name__ == '__main__':
    unittest.main()