tests/test_data_layer.py

"""Data-layer tests. The single vendor call (data.historical._fetch_yf) is
replaced by a recording stand-in, so request windowing, chunking, clamping,
stitching and resampling are exercised without a network round trip. Only
that seam is substituted - every behaviour asserted below is the project's
own logic.

Candle spacing inside each returned window is deliberately coarse to keep
the frames small; what these tests assert is the requested windows and the
stitched/resampled result, not the row spacing.

Run with:  python -m unittest tests.test_data_layer -v
"""
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),  # yfinance ships extra columns
        },
        index=idx,
    )


class _RecordingFetcher:
    """Stands in for the yfinance call, recording every requested window."""

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