tests/test_forecast_api.py

import pytest
from unittest.mock import AsyncMock, MagicMock
from fastapi.testclient import TestClient
import pandas as pd
import app.main as main

# Create TestClient without context manager to bypass actual lifespan events
client = TestClient(main.app, raise_server_exceptions=False)

@pytest.fixture(autouse=True)
def setup_mocks(monkeypatch):
    mock_db = AsyncMock()
    mock_db.connected = True
    
    mock_timesfm = AsyncMock()
    mock_timesfm.model = MagicMock()
    
    mock_data = AsyncMock()
    mock_chart = MagicMock()
    
    monkeypatch.setattr(main, "db_service", mock_db)
    monkeypatch.setattr(main, "timesfm_service", mock_timesfm)
    monkeypatch.setattr(main, "data_service", mock_data)
    monkeypatch.setattr(main, "chart_service", mock_chart)
    
    # Disable rate limiting for testing by mocking the _check method
    monkeypatch.setattr("app.middleware.RateLimitMiddleware._check", lambda self, ip, path: (False, 0))
    
    return mock_db, mock_timesfm, mock_data, mock_chart

def test_health_check_healthy(setup_mocks):
    mock_db, _, _, _ = setup_mocks
    mock_db._execute.return_value = [{"1": 1}]
    
    response = client.get("/health")
    assert response.status_code == 200
    data = response.json()
    assert data["status"] == "healthy"
    assert data["database_connected"] is True
    assert data["model_loaded"] is True

def test_health_check_db_disconnected(setup_mocks):
    mock_db, _, _, _ = setup_mocks
    mock_db.connected = False
    
    response = client.get("/health")
    assert response.status_code == 200
    assert response.json()["database_connected"] is False

def test_get_forecast_cache_hit(setup_mocks):
    mock_db, _, _, _ = setup_mocks
    mock_db.get_cached_forecast.return_value = {
        "symbol": "AAPL",
        "name": "Apple Inc.",
        "exchange": "NASDAQ",
        "currency": "USD",
        "current_price": 150.0,
        "last_updated": "2026-06-19T00:00:00",
        "horizon_days": 20,
        "point_forecast": 155.0,
        "percentage_change": 3.33,
        "quantiles": {"p10": [148.0], "p50": [155.0], "p90": [162.0]},
        "chart_svg": "<svg></svg>",
        "methodology_version": "timesfm-2.5-200m-v1.0"
    }
    
    response = client.get("/api/v1/forecast/AAPL?horizon=20")
    assert response.status_code == 200
    assert response.json()["symbol"] == "AAPL"
    mock_db.get_cached_forecast.assert_called_once_with("AAPL", 20)

def test_get_forecast_cache_miss_and_generate(setup_mocks):
    mock_db, mock_timesfm, mock_data, mock_chart = setup_mocks
    mock_db.get_cached_forecast.return_value = None
    
    mock_df = pd.DataFrame(
        {"Close": [140.0 + i for i in range(100)]},
        index=pd.date_range("2026-01-01", periods=100)
    )
    mock_df.attrs["name"] = "Apple Inc."
    mock_df.attrs["exchange"] = "NASDAQ"
    mock_df.attrs["currency"] = "USD"
    mock_df.attrs["pe_ratio"] = 28.4
    mock_df.attrs["dividend_yield"] = 0.0055
    mock_df.attrs["fifty_two_week_low"] = 130.0
    mock_df.attrs["fifty_two_week_high"] = 200.0
    mock_df.attrs["market_cap"] = 2500000000000.0
    mock_data.get_stock_data.return_value = mock_df
    
    def mock_predict(historical_prices, horizon):
        return {
            "point_forecast": [245.0] * horizon,
            "quantiles": {
                "p10": [240.0] * horizon,
                "p50": [245.0] * horizon,
                "p90": [250.0] * horizon
            }
        }
    mock_timesfm.predict.side_effect = mock_predict
    mock_chart.generate_forecast_chart.return_value = "<svg>chart</svg>"
    
    response = client.get("/api/v1/forecast/AAPL?horizon=20")
    assert response.status_code == 200
    data = response.json()
    assert data["symbol"] == "AAPL"
    assert data["current_price"] == 239.0
    
    # Assert historical prices presence
    assert "historical_prices" in data
    assert isinstance(data["historical_prices"], list)
    assert len(data["historical_prices"]) > 0
    assert "time" in data["historical_prices"][0]
    assert "value" in data["historical_prices"][0]
    
    # Assert backtest accuracy presence and format
    assert "backtest_accuracy" in data
    assert data["backtest_accuracy"] is not None
    assert "accuracy_5d" in data["backtest_accuracy"]
    assert "accuracy_20d" in data["backtest_accuracy"]
    assert "predictions_5d" in data["backtest_accuracy"]
    assert "predictions_20d" in data["backtest_accuracy"]
    
    # Assert financial metrics presence
    assert "financial_metrics" in data
    assert data["financial_metrics"] is not None
    assert data["financial_metrics"]["pe_ratio"] == 28.4
    assert data["financial_metrics"]["dividend_yield"] == 0.0055
    assert data["financial_metrics"]["fifty_two_week_low"] == 130.0
    assert data["financial_metrics"]["fifty_two_week_high"] == 200.0
    assert data["financial_metrics"]["market_cap"] == 2500000000000.0
    
    mock_db.cache_forecast.assert_called_once()