tests/test_oldwang_features.py

import numpy as np
import pandas as pd

from models.predictor import (
    FEATURE_COLUMNS,
    OLDWANG_WEIGHT_OVERLAY,
    _apply_oldwang_weight_overlay,
    _build_features,
    _build_oldwang_context,
)


OLDWANG_FEATURES = [
    "oldwang_triple_bull",
    "oldwang_triple_bear",
    "oldwang_ma5_hold",
    "oldwang_trust_ma10_guard",
    "oldwang_trust_ma10_broken",
    "oldwang_foreign_ma20_guard",
    "oldwang_foreign_ma20_broken",
    "oldwang_volume_spike",
    "oldwang_volume_high_break",
    "oldwang_volume_low_guard",
    "oldwang_volume_low_break",
    "oldwang_gap_guard",
    "oldwang_gap_filled",
    "oldwang_bull_score",
    "oldwang_bear_score",
]


def _oldwang_frame(n: int = 80) -> pd.DataFrame:
    close = pd.Series(np.linspace(100, 145, n))
    volume = np.full(n, 10_000.0)
    volume[50] = 45_000.0
    df = pd.DataFrame(
        {
            "open": close.shift(1).fillna(close.iloc[0]) + 0.2,
            "high": close + 1.5,
            "low": close - 1.0,
            "close": close,
            "volume": volume,
            "ma5": close.rolling(5, min_periods=1).mean(),
            "ma10": close.rolling(10, min_periods=1).mean(),
            "ma20": close.rolling(20, min_periods=1).mean(),
            "ma60": close.rolling(60, min_periods=1).mean(),
            "rsi": 55.0,
            "bb_pct_b": 0.5,
            "k": 50.0,
            "d": 48.0,
            "volume_ratio": pd.Series(volume).rolling(20, min_periods=1).mean() / 10_000.0,
            "atr_ratio": 0.02,
            "obv_trend": 0.01,
            "volatility_20d": 0.02,
            "foreign_net": 1_000.0,
            "trust_net": 500.0,
            "dealer_net": 0.0,
            "institutional_net": 1_500.0,
        }
    )
    df["macd"] = np.linspace(-1, 1, n)
    df["macd_signal"] = np.linspace(-0.8, 0.7, n)
    df["macd_hist"] = df["macd"] - df["macd_signal"]
    return df


def test_oldwang_features_are_production_columns_and_non_null():
    df = _oldwang_frame()
    feat = _build_features(df)

    for col in OLDWANG_FEATURES:
        assert col in FEATURE_COLUMNS
        assert col in feat.columns
        assert feat[col].replace([np.inf, -np.inf], np.nan).notna().all()


def test_oldwang_bull_context_activates_on_uptrend_with_institutional_buying():
    df = _oldwang_frame()
    feat = _build_features(df)
    context = _build_oldwang_context(df, feat)

    assert feat["oldwang_triple_bull"].tail(10).max() == 1.0
    assert feat["oldwang_trust_ma10_guard"].tail(10).max() == 1.0
    assert feat["oldwang_foreign_ma20_guard"].tail(10).max() == 1.0
    assert feat["oldwang_bull_score"].tail(10).max() >= 3.0
    assert context["bull_score"] >= 3.0
    assert any("三陽開泰" in reason for reason in context["bull_reasons"])
    assert any("外資" in reason for reason in context["bull_reasons"])
    assert context["key_levels"]["ma20"] is not None


def test_oldwang_context_reports_volume_and_gap_risks():
    df = _oldwang_frame()
    feat = _build_features(df)
    feat.loc[feat.index[-1], "oldwang_volume_low_break"] = 1.0
    feat.loc[feat.index[-1], "oldwang_gap_filled"] = 1.0
    feat.loc[feat.index[-1], "oldwang_bear_score"] = 2.0

    context = _build_oldwang_context(df, feat)

    assert context["bear_score"] == 2.0
    assert any("爆大量" in reason for reason in context["risk_reasons"])
    assert any("缺口回補" in reason for reason in context["risk_reasons"])


def test_oldwang_weight_overlay_applies_validated_top40_weights():
    buy, sell, meta = _apply_oldwang_weight_overlay(
        buy_prob=0.55,
        sell_prob=0.30,
        oldwang_context={"bull_score": 3, "bear_score": 1},
        config=OLDWANG_WEIGHT_OVERLAY,
    )

    assert buy == 0.66
    assert sell == 0.30
    assert meta["applied"] is True
    assert meta["validation_stock_count"] == 40


def test_oldwang_weight_overlay_applies_validated_bear_gate():
    buy, sell, meta = _apply_oldwang_weight_overlay(
        buy_prob=0.70,
        sell_prob=0.20,
        oldwang_context={"bull_score": 1, "bear_score": 3},
        config=OLDWANG_WEIGHT_OVERLAY,
    )

    assert buy == 0.0
    assert sell == 0.20
    assert meta["applied"] is True


def test_oldwang_weight_overlay_can_be_disabled():
    buy, sell, meta = _apply_oldwang_weight_overlay(
        buy_prob=0.55,
        sell_prob=0.30,
        oldwang_context={"bull_score": 3, "bear_score": 1},
        enabled=False,
    )

    assert buy == 0.55
    assert sell == 0.30
    assert meta["applied"] is False