ai-trade/ai_trade/tests/test_pull_worker.py

"""Tests for pull_worker (Phase 3a additions)."""
from collections import deque
from unittest.mock import AsyncMock, MagicMock

import pytest

from ai_trade.state import PollState


async def test_minute_polling_cycle_updates_state_minute_bars():
    """KIS REST mock 의 분봉 데이터가 state.minute_bars[ticker] deque 에 들어간다."""
    from ai_trade.pull_worker import _run_kis_minute_cycle

    state = PollState()
    state.portfolio = {"holdings": [{"ticker": "005930"}, {"ticker": "000660"}]}
    state.screener_preview = {
        "items": [{"ticker": "005930"}, {"ticker": "035720"}]
    }

    kis_client_mock = MagicMock()
    kis_client_mock.get_minute_ohlcv = AsyncMock(side_effect=[
        [{"datetime": "2026-05-18T09:00:00+09:00", "open": 78000,
          "high": 78500, "low": 77900, "close": 78300, "volume": 12345}],
        [{"datetime": "2026-05-18T09:00:00+09:00", "open": 180000,
          "high": 181000, "low": 179800, "close": 180500, "volume": 5000}],
        [{"datetime": "2026-05-18T09:00:00+09:00", "open": 51000,
          "high": 51200, "low": 50800, "close": 51100, "volume": 8000}],
    ])
    kis_client_mock.get_asking_price = AsyncMock(return_value={
        "bid_total": 600, "ask_total": 400, "bid_ratio": 0.6,
        "current_price": 51100, "as_of": "2026-05-18T09:00:30+09:00",
    })

    await _run_kis_minute_cycle(kis_client_mock, state)

    # 3 unique tickers (005930, 000660, 035720)
    assert "005930" in state.minute_bars
    assert "000660" in state.minute_bars
    assert "035720" in state.minute_bars
    assert len(state.minute_bars["005930"]) >= 1
    # asking_price 만 screener-only ticker (035720) 에 들어가야 함
    # (portfolio = 005930, 000660 는 WebSocket 으로 들어옴)
    assert "035720" in state.asking_price


def test_websocket_message_updates_state_asking_price():
    """WebSocket callback factory → state.asking_price 갱신."""
    from ai_trade.pull_worker import make_asking_price_callback

    state = PollState()
    cb = make_asking_price_callback(state)
    cb("005930", {"bid_total": 1000, "ask_total": 800, "bid_ratio": 0.555,
                  "current_price": 78500, "as_of": "2026-05-18T10:00:00+09:00"})
    assert state.asking_price["005930"]["bid_total"] == 1000
    assert "asking_price/005930" in state.last_updated


async def test_post_close_cycle_updates_chronos_predictions():
    """mock kis + mock chronos → state.chronos_predictions + state.daily_ohlcv 갱신."""
    from unittest.mock import AsyncMock, MagicMock
    from ai_trade.pull_worker import _run_post_close_cycle
    from ai_trade.chronos_predictor import ChronosPrediction
    from ai_trade.state import PollState

    state = PollState()
    state.portfolio = {"holdings": [{"ticker": "005930"}]}
    state.screener_preview = {"items": [{"ticker": "000660"}]}

    kis_mock = MagicMock()
    daily_005930 = [{"datetime": f"2026-05-{i+1:02d}", "open": 100, "high": 105,
                     "low": 95, "close": 100 + i, "volume": 1000} for i in range(60)]
    daily_000660 = [{"datetime": f"2026-05-{i+1:02d}", "open": 200, "high": 210,
                     "low": 190, "close": 200 + i, "volume": 2000} for i in range(60)]
    # _run_post_close_cycle iterates tickers and calls get_daily_ohlcv per ticker.
    # Order depends on set() so use side_effect mapping if possible, otherwise list.
    async def fake_daily(ticker, days=60):
        if ticker == "005930":
            return daily_005930
        if ticker == "000660":
            return daily_000660
        return []
    kis_mock.get_daily_ohlcv = AsyncMock(side_effect=fake_daily)

    chronos_mock = MagicMock()
    chronos_mock.predict_batch = MagicMock(return_value={
        "005930": ChronosPrediction(0.02, -0.01, 0.04, 0.85, "2026-05-18T16:00:00+09:00"),
        "000660": ChronosPrediction(0.03, -0.02, 0.06, 0.75, "2026-05-18T16:00:00+09:00"),
    })

    await _run_post_close_cycle(kis_mock, chronos_mock, state)

    assert "005930" in state.chronos_predictions
    assert "000660" in state.chronos_predictions
    assert state.chronos_predictions["005930"]["median"] == 0.02
    assert state.chronos_predictions["005930"]["conf"] == 0.85
    assert "005930" in state.daily_ohlcv
    assert "chronos/005930" in state.last_updated


def test_poll_loop_calls_generate_signals_after_cycle(monkeypatch):
    """Phase 4: generate_signals 가 cycle 후 state.signals 를 갱신한다."""
    from unittest.mock import MagicMock
    from ai_trade.state import PollState
    from ai_trade.signal_generator import generate_signals

    state = PollState()
    state.portfolio = {"holdings": [{
        "ticker": "005930", "name": "삼성전자",
        "avg_price": 75000, "current_price": 69000,
        "pnl_pct": -0.08, "profit_rate": -8.0,
        "quantity": 100, "broker": "키움",
    }]}
    state.screener_preview = {"items": []}

    dedup = MagicMock()
    dedup.is_recent.return_value = False

    settings = MagicMock()
    settings.stop_loss_pct = -0.07
    settings.take_profit_pct = 0.15
    settings.chronos_spread_threshold = 0.6
    settings.asking_bid_ratio_threshold = 0.6
    settings.confidence_threshold = 0.7
    settings.min_momentum_for_buy = "strong_up"

    generate_signals(state, dedup, settings)

    assert "005930" in state.signals
    assert state.signals["005930"]["action"] == "sell"
    assert state.signals["005930"]["confidence_webai"] == 1.0
    dedup.record.assert_called_with("005930", "sell", confidence=1.0)


async def test_post_close_fires_at_1601_when_not_yet_today(monkeypatch):
    """F3 — 16:01에 깬 cycle도 오늘 post_close 안 돌렸으면 호출됨 (회귀 방지)."""
    from datetime import datetime as _dt
    from zoneinfo import ZoneInfo as _ZI
    import asyncio as _asyncio

    from ai_trade import pull_worker

    _kst = _ZI("Asia/Seoul")
    now_at_1601 = _dt(2026, 5, 18, 16, 1, tzinfo=_kst)

    class FrozenDateTime:
        @staticmethod
        def now(tz=None):
            return now_at_1601

    monkeypatch.setattr(pull_worker, "datetime", FrozenDateTime)
    monkeypatch.setattr(pull_worker, "_is_market_day", lambda n: True)
    monkeypatch.setattr(pull_worker, "_is_polling_window", lambda n: True)
    monkeypatch.setattr(pull_worker, "_next_interval", lambda n: 0.01)
    monkeypatch.setattr(pull_worker, "_run_polling_cycle", AsyncMock())
    monkeypatch.setattr(pull_worker, "update_minute_momentum_for_all", lambda s: None)
    post_close = AsyncMock()
    monkeypatch.setattr(pull_worker, "_run_post_close_cycle", post_close)

    state = MagicMock()
    chronos = MagicMock()
    kis = MagicMock()
    shutdown = _asyncio.Event()

    async def _stop_soon():
        await _asyncio.sleep(0.05)
        shutdown.set()

    _asyncio.create_task(_stop_soon())
    await pull_worker.poll_loop(
        client=MagicMock(),
        state=state,
        shutdown=shutdown,
        kis_client=kis,
        chronos=chronos,
        dedup=None,
        settings=None,
    )

    assert post_close.await_count >= 1, "post-close가 16:01에 호출되지 않음 (F3 회귀)"