tests/test_predictor.py

"""
GranitePredictor 단위 테스트.

실제 모델 다운로드 없이 statsforecast AutoARIMA fallback 경로를 검증합니다.
"""

from __future__ import annotations

import numpy as np
import pandas as pd
import pytest

from src.analytics.predictor import ForecastResult, GranitePredictor, _monthly_to_yearly


class TestMonthlyToYearly:
    def test_last_month_aggregation(self) -> None:
        monthly = np.arange(24, dtype=float)  # 24개월 = 2년
        yearly = _monthly_to_yearly(monthly, agg="last")
        assert len(yearly) == 2
        assert yearly[0] == 11.0   # 첫 해의 마지막 달
        assert yearly[1] == 23.0   # 두 번째 해의 마지막 달

    def test_mean_aggregation(self) -> None:
        monthly = np.ones(12, dtype=float) * 5.0  # 12개월, 모두 5
        yearly = _monthly_to_yearly(monthly, agg="mean")
        assert len(yearly) == 1
        assert yearly[0] == pytest.approx(5.0)

    def test_remainder_months_handled(self) -> None:
        monthly = np.ones(14, dtype=float)  # 14개월 = 1년 + 2개월 나머지
        yearly = _monthly_to_yearly(monthly)
        assert len(yearly) == 2  # 1년 + 나머지


class TestForecastResult:
    def test_to_dict_has_all_keys(self) -> None:
        result = ForecastResult(
            schul_code="TEST",
            target_col="student_count",
            forecast_years=[2024, 2025],
            point_forecast=[80.0, 70.0],
            lower_bound=[60.0, 50.0],
            upper_bound=[100.0, 90.0],
            model_version="arima-fallback",
            context_years=[2020, 2021, 2022, 2023],
            context_values=[120.0, 110.0, 100.0, 90.0],
        )
        d = result.to_dict()
        required = {"schul_code", "target_col", "forecast_years", "point_forecast",
                    "lower_bound", "upper_bound", "model_version"}
        assert required.issubset(d.keys())


class TestGranitePredictor:
    def _make_timeseries(self, n: int = 5) -> pd.Series:
        years = list(range(2019, 2019 + n))
        values = [120.0, 110.0, 98.0, 87.0, 75.0][:n]
        return pd.Series(dict(zip(years, values)))

    def test_insufficient_data_raises(self) -> None:
        predictor = GranitePredictor()
        ts = pd.Series({2023: 50.0})  # 단 1개 → 부족
        with pytest.raises(ValueError, match="부족"):
            predictor.predict("TEST", ts)

    def test_empty_series_raises(self) -> None:
        predictor = GranitePredictor()
        with pytest.raises(ValueError):
            predictor.predict("TEST", pd.Series(dtype=float))

    def test_arima_fallback_produces_result(self) -> None:
        """TTM 로드 없이 ARIMA fallback 이 결과를 반환하는지 검증."""
        predictor = GranitePredictor()
        # TTM 로드를 강제로 실패시켜 ARIMA fallback 으로 이동
        predictor._model_version = predictor._MODEL_VERSION_ARIMA
        ts = self._make_timeseries()

        try:
            result = predictor.predict("TEST", ts, horizon_years=3)
            assert isinstance(result, ForecastResult)
            assert len(result.forecast_years) == 3
            assert len(result.point_forecast) == 3
            assert all(v >= 0 for v in result.point_forecast)
            assert result.model_version == predictor._MODEL_VERSION_ARIMA
        except ImportError:
            pytest.skip("statsforecast 미설치, ARIMA fallback 건너뜀")

    def test_forecast_years_sequential(self) -> None:
        """예측 연도가 순차적으로 증가하는지 검증."""
        predictor = GranitePredictor()
        predictor._model_version = predictor._MODEL_VERSION_ARIMA
        ts = self._make_timeseries()

        try:
            result = predictor.predict("TEST", ts, horizon_years=3)
            years = result.forecast_years
            for i in range(1, len(years)):
                assert years[i] == years[i - 1] + 1
        except ImportError:
            pytest.skip("statsforecast 미설치")

    def test_lower_bound_le_upper_bound(self) -> None:
        predictor = GranitePredictor()
        predictor._model_version = predictor._MODEL_VERSION_ARIMA
        ts = self._make_timeseries()

        try:
            result = predictor.predict("TEST", ts, horizon_years=3)
            for lo, hi in zip(result.lower_bound, result.upper_bound):
                assert lo <= hi
        except ImportError:
            pytest.skip("statsforecast 미설치")