"""Tests for ChronosPredictor (model mock).""" from unittest.mock import MagicMock, patch import numpy as np import pytest @pytest.fixture def mock_pipeline(): """Mock BaseChronosPipeline.from_pretrained returning a mock pipeline object.""" with patch("chronos.BaseChronosPipeline") as cls: cls.__name__ = "BaseChronosPipeline" instance = MagicMock() # ChronosBolt API: predict_quantiles returns (quantiles_tensor, mean_tensor) # Modern (predict_quantiles) branch will be used since hasattr(MagicMock, "predict_quantiles") is True. cls.from_pretrained.return_value = instance yield instance @pytest.fixture def mock_torch_cpu(): with patch("torch.cuda.is_available", return_value=False): yield def _daily_ohlcv(close_seq): return [{"datetime": f"2026-05-{i+1:02d}", "open": c, "high": c, "low": c, "close": c, "volume": 1000} for i, c in enumerate(close_seq)] def _mk_quantiles_tensor(q10_price: float, q50_price: float, q90_price: float): """Helper: build predict_quantiles return tensor shape [1, 1, 3].""" import torch return torch.tensor([[[q10_price, q50_price, q90_price]]], dtype=torch.float32) def test_predict_batch_returns_prediction_dict(mock_pipeline, mock_torch_cpu): """mock predict_quantiles → dict[ticker, ChronosPrediction]. last_close=100, q50=102 → median≈+2%.""" quantiles = _mk_quantiles_tensor(101.5, 102.0, 102.5) # narrow around 102 mock_pipeline.predict_quantiles.return_value = (quantiles, None) from signal_v2.chronos_predictor import ChronosPredictor, ChronosPrediction predictor = ChronosPredictor(model_name="mock-model") daily = {"005930": _daily_ohlcv([100] * 60)} result = predictor.predict_batch(daily) assert "005930" in result pred = result["005930"] assert isinstance(pred, ChronosPrediction) assert abs(pred.median - 0.02) < 0.001 def test_conf_high_when_distribution_narrow(mock_pipeline, mock_torch_cpu): """좁은 distribution (q90-q10 작음, median 0 아님) → conf ≈ 1.""" # last_close=100, q10=101.99, q50=102.00, q90=102.01 # returns: q10=0.0199, q50=0.02, q90=0.0201 # spread = (0.0201 - 0.0199) / max(0.02, 0.001) = 0.0002/0.02 = 0.01 → conf = 1 - 0.005 = 0.995 quantiles = _mk_quantiles_tensor(101.99, 102.0, 102.01) mock_pipeline.predict_quantiles.return_value = (quantiles, None) from signal_v2.chronos_predictor import ChronosPredictor predictor = ChronosPredictor(model_name="mock-model") daily = {"005930": _daily_ohlcv([100] * 60)} result = predictor.predict_batch(daily) assert result["005930"].conf > 0.8 def test_conf_low_when_distribution_wide(mock_pipeline, mock_torch_cpu): """넓은 distribution → conf ≈ 0.""" # last_close=100, q10=70, q50=100, q90=130 # returns: q10=-0.3, q50=0.0, q90=0.3 # spread = (0.3 - (-0.3)) / max(0.0, 0.001) = 0.6 / 0.001 = 600 → conf = max(0, 1 - 300) = 0 quantiles = _mk_quantiles_tensor(70.0, 100.0, 130.0) mock_pipeline.predict_quantiles.return_value = (quantiles, None) from signal_v2.chronos_predictor import ChronosPredictor predictor = ChronosPredictor(model_name="mock-model") daily = {"005930": _daily_ohlcv([100] * 60)} result = predictor.predict_batch(daily) assert result["005930"].conf < 0.3 def test_return_computed_from_price_relative_to_last_close(mock_pipeline, mock_torch_cpu): """price 예측 → last_close 대비 return 변환. last_close=100, q50=110 → return ≈ +10%.""" quantiles = _mk_quantiles_tensor(109.0, 110.0, 111.0) mock_pipeline.predict_quantiles.return_value = (quantiles, None) from signal_v2.chronos_predictor import ChronosPredictor predictor = ChronosPredictor(model_name="mock-model") # last close = 100 daily = {"005930": _daily_ohlcv(list(range(41, 101)))} # last = 100 result = predictor.predict_batch(daily) assert abs(result["005930"].median - 0.10) < 0.001