ai-trade/ai_trade/tests/test_momentum_classifier.py

"""Tests for minute momentum classifier."""
from collections import deque

from ai_trade.momentum_classifier import (
    aggregate_1min_to_5min, classify_minute_momentum,
    STRONG_UP, WEAK_UP, NEUTRAL, WEAK_DOWN, STRONG_DOWN,
)


def _bar(open_, high, low, close, volume):
    return {
        "datetime": "2026-05-18T09:00:00+09:00",
        "open": open_, "high": high, "low": low, "close": close, "volume": volume,
    }


def _make_chunks(num_chunks_up: int, num_chunks_total: int, base_vol: int = 1000):
    """num_chunks_total 개의 5-bar 청크. num_chunks_up 청크는 양봉, 나머지는 음봉.
    각 청크는 5개 1분봉. 거래량 = base_vol per bar.
    """
    bars = []
    for i in range(num_chunks_total):
        is_up = i < num_chunks_up
        o, c = (100, 110) if is_up else (110, 100)
        for j in range(5):
            bars.append(_bar(o, max(o, c) + 5, min(o, c) - 5, c, base_vol))
    return bars


def test_strong_up_5_consecutive_green_with_high_volume():
    """직전 5개 5분봉 모두 양봉 + 거래량 1.5x → STRONG_UP."""
    # 60분 (12 5분봉) 데이터: 7 normal + 5 high-vol up
    older = _make_chunks(num_chunks_up=3, num_chunks_total=7, base_vol=1000)
    recent = _make_chunks(num_chunks_up=5, num_chunks_total=5, base_vol=2500)
    minute_bars = deque(older + recent, maxlen=60)
    assert classify_minute_momentum(minute_bars) == STRONG_UP


def test_weak_up_3of5_green_normal_volume():
    """직전 5개 5분봉 중 3-4개 양봉 + 거래량 ≥ 1.0x → WEAK_UP."""
    older = _make_chunks(num_chunks_up=3, num_chunks_total=7, base_vol=1000)
    # 5 chunks: 3 up + 2 down, normal vol
    recent_up = _make_chunks(num_chunks_up=3, num_chunks_total=3, base_vol=1000)
    recent_down = _make_chunks(num_chunks_up=0, num_chunks_total=2, base_vol=1000)
    minute_bars = deque(older + recent_up + recent_down, maxlen=60)
    assert classify_minute_momentum(minute_bars) == WEAK_UP


def test_neutral_mixed():
    """up_count=2, vol normal → NEUTRAL (rule 미해당)."""
    older = _make_chunks(num_chunks_up=3, num_chunks_total=7, base_vol=1000)
    recent_up = _make_chunks(num_chunks_up=2, num_chunks_total=2, base_vol=1000)
    recent_down = _make_chunks(num_chunks_up=0, num_chunks_total=3, base_vol=1000)
    minute_bars = deque(older + recent_up + recent_down, maxlen=60)
    # up_count=2, vol_mult=1.0 → 어느 분기 조건도 만족 안 함 → NEUTRAL
    assert classify_minute_momentum(minute_bars) == NEUTRAL


def test_weak_down_low_green_low_volume():
    """up_count <= 2 + vol < 1.0 → WEAK_DOWN."""
    older = _make_chunks(num_chunks_up=3, num_chunks_total=7, base_vol=1000)
    recent_up = _make_chunks(num_chunks_up=1, num_chunks_total=1, base_vol=500)
    recent_down = _make_chunks(num_chunks_up=0, num_chunks_total=4, base_vol=500)
    minute_bars = deque(older + recent_up + recent_down, maxlen=60)
    # recent 5 chunks avg vol = 500, long 12 avg ≈ (7*1000 + 5*500) / 12 ≈ 791 → vol_mult ≈ 0.63
    assert classify_minute_momentum(minute_bars) == WEAK_DOWN


def test_strong_down_5_consecutive_red_high_volume():
    """직전 5개 5분봉 모두 음봉 + 거래량 1.5x → STRONG_DOWN."""
    older = _make_chunks(num_chunks_up=3, num_chunks_total=7, base_vol=1000)
    recent = _make_chunks(num_chunks_up=0, num_chunks_total=5, base_vol=2500)
    minute_bars = deque(older + recent, maxlen=60)
    assert classify_minute_momentum(minute_bars) == STRONG_DOWN


def test_aggregate_1min_to_5min_correctness():
    """5 1분봉 → 1개 5분봉 — open/close/high/low/volume 정확."""
    bars = [
        _bar(100, 105, 99, 102, 1000),
        _bar(102, 108, 101, 107, 1500),
        _bar(107, 110, 105, 106, 800),
        _bar(106, 109, 104, 108, 1200),
        _bar(108, 112, 107, 111, 900),
    ]
    result = aggregate_1min_to_5min(bars)
    assert len(result) == 1
    assert result[0]["open"] == 100  # 첫 bar
    assert result[0]["close"] == 111  # 마지막 bar
    assert result[0]["high"] == 112  # max
    assert result[0]["low"] == 99    # min
    assert result[0]["volume"] == 5400  # sum