Haocode/tests/verify_math_render.py

# -*- coding: utf-8 -*-
"""公式渲染端到端验证：真实会话消息 → 真实前端管线 → DOM 断言

运行: PYTHONIOENCODING=utf-8 QT_QPA_PLATFORM=offscreen python tests/verify_math_render.py

铁律：不污染真实 DB —— 先把 data/chat_history.db 复制到临时文件，再指向副本。
"""
import os
import sys
import json
import time
import shutil
import tempfile

sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
os.environ.setdefault("QT_QPA_PLATFORM", "offscreen")
os.environ.setdefault("HAOCODE_RENDER", "software")
os.environ["QTWEBENGINE_CHROMIUM_FLAGS"] = "--disable-gpu"

_REAL_DB = os.path.join(os.path.dirname(__file__), "..", "data", "chat_history.db")
_TMP_DB = os.path.join(tempfile.gettempdir(), f"haocode_verify_math_{os.getpid()}.db")
shutil.copyfile(_REAL_DB, _TMP_DB)

import core.db_manager as _dbm  # noqa: E402
_dbm._DEFAULT_DB = _TMP_DB

from PyQt6.QtWidgets import QApplication  # noqa: E402
from PyQt6.QtCore import QTimer  # noqa: E402
from ui.views.main_window import MainWindow  # noqa: E402

app = QApplication(sys.argv)
results = []
win = {"w": None, "sid": None}

JS_QUERY = r"""(function(){
    var scope = document.getElementById('chat-container') || document.body;
    var q = function(s){ return scope.querySelectorAll(s).length; };
    // 只取助手消息（user 消息按设计是纯文本，本来就不渲染公式）
    var asst = scope.querySelectorAll('.message-wrapper.assistant');
    var asstScope = asst.length ? asst[asst.length - 1] : scope;
    // 可见文本：剔除 katex-mathml（CSS 视觉隐藏的 MathML 层，内部保留原始 tex 属正常）
    var clone = asstScope.cloneNode(true);
    clone.querySelectorAll('.katex-mathml').forEach(function(e){ e.parentNode.removeChild(e); });
    var visTxt = clone.textContent || '';
    var codeTxt = '';
    asstScope.querySelectorAll('pre, code').forEach(function(e){ codeTxt += e.textContent + '\n'; });
    return {
        katex: q('.katex'),
        display: q('.katex-display'),
        mathml: q('.katex-mathml'),
        vislayer: q('.katex-html'),
        leakedPlaceholder: visTxt.indexOf('@@K') !== -1,
        rawBackslash: visTxt.indexOf('\\operatorname') !== -1,
        bareBracketFormula: visTxt.indexOf('P_4=\\operatorname') !== -1,
        codeHasDollar: codeTxt.indexOf('$x + y$') !== -1,
        codeHasArr: codeTxt.indexOf('arr[0]') !== -1
    };
})()"""


def check(name, cond, extra=""):
    print(("  PASS  " if cond else "  FAIL  ") + name + (("  | " + str(extra)) if extra else ""))
    results.append(bool(cond))


def run_js(js, timeout_s=15):
    r = {"v": None, "d": False}

    def cb(val):
        if isinstance(val, str):
            try:
                val = json.loads(val)
            except Exception:
                pass
        r["v"] = val
        r["d"] = True

    b = win["w"].browser
    if hasattr(b, "execute_js_async"):
        b.execute_js_async(js, cb)
    else:
        b.page().runJavaScript(js, cb)
    t0 = time.time()
    while not r["d"] and time.time() - t0 < timeout_s:
        app.processEvents()
        time.sleep(0.05)
    return r["v"]


def boot():
    w = MainWindow()
    win["w"] = w
    sid = None
    for s in w.db.get_all_sessions():
        if s.get("title") == "公式渲染验收":
            sid = s["id"]
            break
    if not sid:
        print("FAIL 未找到「公式渲染验收」会话（先跑 tests/inject_math_demo.py）")
        app.quit()
        return
    win["sid"] = sid
    print(f"会话: {sid}")
    w.load_messages_to_web(sid)
    QTimer.singleShot(9000, phase_check)


def phase_check():
    for _ in range(40):
        if run_js("window.jsReady === true ? 1 : 0", timeout_s=3) == 1:
            break
        time.sleep(0.5)

    r = run_js(JS_QUERY)
    if not isinstance(r, dict):
        print(f"FAIL DOM 查询失败: {r!r}")
        app.quit()
        return
    print("\n  DOM 统计: " + json.dumps(r, ensure_ascii=False))
    check("KaTeX 渲染出公式（.katex > 0）", r["katex"] > 0, f"katex={r['katex']}")
    check("块公式 7 个（.katex-display == 7）", r["display"] == 7, f"display={r['display']}")
    check("可视层存在（.katex-html > 0）", r["vislayer"] > 0, f"vislayer={r['vislayer']}")
    check("无占位符泄漏", not r["leakedPlaceholder"])
    check("可见文本无原始 tex 残留", not r["rawBackslash"])
    check("无裸括号公式残留", not r["bareBracketFormula"])
    check("代码块内 $x + y$ 保持原样", r["codeHasDollar"])
    check("代码块内 arr[0] 保持原样", r["codeHasArr"])

    try:
        shot_path = os.path.join(os.path.dirname(__file__), "_tmp_math_render.png")
        win["w"].browser.grab().save(shot_path)
        print(f"  截图: {shot_path}")
    except Exception as e:
        print(f"  截图失败: {e}")

    print(f"\n===== {'ALL PASS' if all(results) else 'HAS FAILURES'}: {sum(results)}/{len(results)} =====")
    try:
        win["w"].close()
    except Exception:
        pass
    app.quit()


QTimer.singleShot(400, boot)
QTimer.singleShot(120000, app.quit)
app.exec()
try:
    os.remove(_TMP_DB)
except Exception:
    pass
sys.exit(0 if all(results) else 1)