kis_trader/_test_kiwoom_condition_realtime.py

#!/usr/bin/env python3
"""
_test_kiwoom_condition_realtime.py — 키움 웹소켓 '실시간 조건검색'(ka10173 / CNSRREQ) 수신 테스트
====================================================================================================
[목적]
키움 신형 WS 로 조건검색식 'momentum' 을 실시간(search_type=1) 등록해,
① 초기 매칭 종목 리스트와 ② 실시간 편입/이탈(REAL, 843=삽입 I/삭제 D, 9001=종목코드)을
수신할 수 있는지 확인한다. (유니버스 실매 전환의 핵심 전제)

[흐름]
1) LOGIN → CNSRLST 로 'momentum' 조건식 seq 를 이름으로 해결
2) CNSRREQ(search_type=1, stex_tp=K) 로 실시간 등록 → 초기 응답(현재 매칭) 출력
3) WATCH_SEC 동안 REAL push(편입/이탈) 수신·출력
4) CNSRCLR 로 실시간 해제 후 종료

[주의]
- 읽기 전용(조회/실시간 수신)만 한다. 주문·저장은 하지 않는다.
- 장 마감 후에는 실시간 편입/이탈 push 가 없을 수 있다(초기 매칭 리스트로 요청 성공만 확인).
- 응답 원문(raw)을 그대로 찍어 필드명이 문서와 달라도 눈으로 확인 가능.
"""
import json
import os
import sys
import threading
import time

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from kis_trader.utils.env import get_env_from_db, get_env_int
from kis_trader.ws.kis_ws import _get_kiwoom_token_cached

TARGET_NAME = "momentum"          # 실시간 등록할 조건식 이름
WATCH_SEC = get_env_int("KIWOOM_COND_RT_WATCH_SEC", 30)  # 실시간 수신 관찰 시간(초)


def _load_kiwoom_creds():
    kw_mock_raw = (get_env_from_db("KIWOOM_MOCK", "") or "").strip().lower()
    is_mock = kw_mock_raw in ("1", "true", "y", "yes", "on")
    if is_mock:
        key = (get_env_from_db("KIWOOM_APP_KEY_MOCK", "") or "").strip()
        secret = (get_env_from_db("KIWOOM_APP_SECRET_MOCK", "") or "").strip()
    else:
        key = (get_env_from_db("KIWOOM_APP_KEY_REAL", "") or "").strip()
        secret = (get_env_from_db("KIWOOM_APP_SECRET_REAL", "") or "").strip()
    if not key or not secret:
        key = key or (get_env_from_db("KIWOOM_APP_KEY", "") or "").strip()
        secret = secret or (get_env_from_db("KIWOOM_APP_SECRET", "") or "").strip()
    return key, secret, is_mock


def _seq_name(item):
    """CNSRLST data 항목: [seq, name] 배열 또는 {seq,name} dict 모두 허용."""
    if isinstance(item, (list, tuple)):
        return (str(item[0]) if len(item) > 0 else ""), (str(item[1]) if len(item) > 1 else "")
    return str(item.get("seq")), str(item.get("name"))


def _extract_code(item):
    """조건검색 결과 항목에서 종목코드 추출 (9001 우선, jmcode 폴백)."""
    if isinstance(item, dict):
        return str(item.get("9001") or item.get("jmcode") or "").strip()
    if isinstance(item, (list, tuple)) and item:
        return str(item[0]).strip()
    return ""


def main() -> int:
    try:
        import websocket  # websocket-client
    except Exception as e:
        print(f"❌ websocket-client 미설치: {e}")
        return 1

    key, secret, is_mock = _load_kiwoom_creds()
    if not key or not secret:
        print("❌ 키움 앱키/시크릿 미설정 — env_config 확인")
        return 1
    token = _get_kiwoom_token_cached(key, secret, is_mock)
    if not token:
        print("❌ 키움 토큰 발급 실패")
        return 1
    print(f"🎫 토큰 OK (mock={is_mock}, 앞8자={token[:8]}…)")

    url = (
        get_env_from_db("KIWOOM_WS_URL_MOCK", "wss://mockapi.kiwoom.com:10000/api/dostk/websocket")
        if is_mock else
        get_env_from_db("KIWOOM_WS_URL_REAL", "wss://api.kiwoom.com:10000/api/dostk/websocket")
    )
    print(f"🌐 WS 연결: {url}")

    st = {
        "seq": None, "registered": False, "cleared": False,
        "reg_deadline": None, "error": None, "close": False,
        "real_events": 0, "initial_codes": None,
    }

    def on_open(ws):
        ws.send(json.dumps({"trnm": "LOGIN", "token": token}))
        print("📡 LOGIN 발송")

    def on_message(ws, message):
        try:
            data = json.loads(message)
        except Exception:
            print(f"📥 (raw) {message[:200]}")
            return
        trnm = data.get("trnm")

        if trnm == "PING":
            try:
                ws.send(message)
            except Exception:
                pass
            return

        if trnm == "LOGIN":
            if str(data.get("return_code")) in ("0", "0.0"):
                print("✅ LOGIN OK → CNSRLST(목록조회)")
                ws.send(json.dumps({"trnm": "CNSRLST"}))
            else:
                st["error"] = f"LOGIN 실패: {data.get('return_msg')}"
                ws.close()
            return

        if trnm == "CNSRLST":
            found = None
            for it in (data.get("data") or []):
                seq, name = _seq_name(it)
                if name.strip().lower() == TARGET_NAME:
                    found = seq
                    break
            if found is None:
                st["error"] = f"'{TARGET_NAME}' 조건식 없음 (목록: {data.get('data')})"
                ws.close()
                return
            st["seq"] = found
            print(f"🎯 '{TARGET_NAME}' seq={found} → CNSRREQ 실시간(search_type=1) 등록")
            ws.send(json.dumps({
                "trnm": "CNSRREQ", "seq": found, "search_type": "1", "stex_tp": "K",
            }))
            return

        if trnm == "CNSRREQ":
            rc = str(data.get("return_code"))
            print("\n===== CNSRREQ 초기 응답 원문 =====")
            print(json.dumps(data, ensure_ascii=False, indent=2)[:2000])
            if rc not in ("0", "0.0"):
                st["error"] = f"CNSRREQ 실패 rc={rc} msg={data.get('return_msg')}"
                ws.close()
                return
            codes = [_extract_code(it) for it in (data.get("data") or [])]
            codes = [c for c in codes if c]
            st["initial_codes"] = codes
            st["registered"] = True
            st["reg_deadline"] = time.time() + WATCH_SEC
            print(f"\n✅ 실시간 등록 성공 — 초기 매칭 {len(codes)}종목: {codes[:20]}")
            print(f"⏳ {WATCH_SEC}초간 실시간 편입/이탈(REAL) 수신 대기…")
            return

        if trnm == "REAL":
            for it in (data.get("data") or []):
                vals = it.get("values") if isinstance(it, dict) else None
                if not isinstance(vals, dict):
                    continue
                code = str(vals.get("9001") or "").strip()
                ins_del = str(vals.get("843") or "").strip()  # I=삽입(편입), D=삭제(이탈)
                sig = str(vals.get("841") or "").strip()
                tm = str(vals.get("20") or "").strip()
                kind = "편입(I)" if ins_del == "I" else ("이탈(D)" if ins_del == "D" else ins_del)
                st["real_events"] += 1
                print(f"   📶 REAL {kind} 종목={code} 신호seq={sig} 시각={tm}")
            return

        if trnm == "CNSRCLR":
            print(f"🧹 CNSRCLR(실시간 해제) 응답: rc={data.get('return_code')}")
            st["cleared"] = True
            ws.close()
            return

        print(f"📥 기타: {json.dumps(data, ensure_ascii=False)[:300]}")

    def on_error(ws, err):
        st["error"] = f"WS 오류: {err}"

    def on_close(ws, code, msg):
        st["close"] = True

    ws = websocket.WebSocketApp(
        url, on_open=on_open, on_message=on_message,
        on_error=on_error, on_close=on_close,
    )
    th = threading.Thread(target=ws.run_forever, kwargs={"ping_interval": 0}, daemon=True)
    th.start()

    # 등록 후 WATCH_SEC 경과 → CNSRCLR 해제 → 종료
    hard_deadline = time.time() + WATCH_SEC + 30
    while not st["close"] and time.time() < hard_deadline:
        time.sleep(0.2)
        if st["error"]:
            break
        if (st["registered"] and not st["cleared"]
                and st["reg_deadline"] and time.time() >= st["reg_deadline"]):
            try:
                ws.send(json.dumps({"trnm": "CNSRCLR", "seq": st["seq"]}))
                print("📤 CNSRCLR(실시간 해제) 발송")
            except Exception:
                pass
            st["reg_deadline"] = None  # 1회만
    try:
        ws.close()
    except Exception:
        pass

    if st["error"]:
        print(f"\n🚨 {st['error']}")
        return 2
    if not st["registered"]:
        print("\n⏱️  실시간 등록 미완료(응답 없음/타임아웃)")
        return 3

    print(f"\n📊 요약: 초기매칭 {len(st['initial_codes'] or [])}종목 · REAL 이벤트 {st['real_events']}건")
    print("🎉 실시간 조건검색 요청/등록/해제 경로 검증 완료"
          + (" (장중 아니라 편입/이탈 push 는 0건일 수 있음)" if st["real_events"] == 0 else ""))
    return 0


if __name__ == "__main__":
    sys.exit(main())