Yb-QaAutomationHub/tests/core/playwright_support.py

"""Playwright 支持 — 浏览器 Page 工厂 + 快速关闭 + 登录态加载 + 失败截图。

快速关闭背景: 本机 Windows 环境下，Playwright 关闭浏览器或退出驱动时会先做
最长约 30s 的优雅关闭等待（等 Chrome 进程自然退出），导致脚本/用例结束慢。
fast_shutdown 强杀浏览器进程树 + node 驱动，让 p.stop() 立即返回。
"""

import os
import subprocess
from contextlib import contextmanager
from datetime import datetime
from pathlib import Path

import pytest

from tests.core.allure_support import allure, attach_png, attach_text
from tests.core.paths import SCREENSHOTS_DIR

# 本机 %LOCALAPPDATA%\ms-playwright 下 Chrome for Testing 存在 SxS 激活失败问题，
# 浏览器迁移到了自定义目录；仅当该目录存在时兜底设置，不影响其他机器。
_CUSTOM_BROWSERS_PATH = "C:/Users/23184/pw-playwright"
if os.path.isdir(_CUSTOM_BROWSERS_PATH):
    os.environ.setdefault("PLAYWRIGHT_BROWSERS_PATH", _CUSTOM_BROWSERS_PATH)


# ═══════════════════════════════════════════════════════════════════
# Page 工厂
# ═══════════════════════════════════════════════════════════════════

@contextmanager
def admin_browser_page(admin_config: dict, browser_name: str = "chromium",
                       role: str | None = None):
    """创建管理端 Playwright Page（skip_login 时加载指定角色的登录态）。

    Yields:
        (page, loaded_role): loaded_role 为实际加载登录态的角色，未加载为 None

    退出时自动 fast_shutdown，跳过本机 ~30s 优雅关闭等待。
    """
    try:
        from playwright.sync_api import sync_playwright
    except ImportError:
        pytest.skip("Playwright 未安装。")

    viewport = admin_config.get("viewport", {"width": 1920, "height": 1080})
    timeout = admin_config.get("timeout_seconds", 120) * 1000
    storage_state_dir = Path(admin_config.get("storage_state_dir", "tests/web/auth"))
    skip_login = admin_config.get("skip_login", True)

    with sync_playwright() as p:
        browser = getattr(p, browser_name).launch(headless=True)

        context_kwargs = {
            "viewport": {"width": viewport["width"], "height": viewport["height"]},
            "locale": "zh-CN",
        }

        loaded_role = None
        if skip_login and role:
            auth_file = storage_state_dir / f"{role}_{browser_name}.json"
            if auth_file.exists():
                context_kwargs["storage_state"] = str(auth_file)
                loaded_role = role

        context = browser.new_context(**context_kwargs)
        page = context.new_page()
        page.set_default_timeout(timeout)

        try:
            yield page, loaded_role
        finally:
            fast_shutdown(p)


def pick_available_role(storage_state_dir: Path, browser_name: str) -> str | None:
    """选择可用的登录态角色：优先 super_admin，否则按文件名排序取第一个。"""
    if not storage_state_dir.exists():
        return None
    files = sorted(storage_state_dir.glob(f"*_{browser_name}.json"))
    if not files:
        return None
    for f in files:
        if f.stem == f"super_admin_{browser_name}":
            return "super_admin"
    return files[0].stem[: -len(f"_{browser_name}")]


def save_web_failure_screenshot(page, request, tag: str) -> str | None:
    """用例失败时整页截图 + DOM 快照（存盘/attach Allure），返回路径；未失败返回 None。

    Args:
        page: Playwright Page
        request: pytest request（依赖根 conftest 的 makereport hook 提供 rep_call）
        tag: 文件名标识，如 "chromium" / "admin_chromium"
    """
    if not (hasattr(request.node, "rep_call") and request.node.rep_call.failed):
        return None
    scenario_dir = SCREENSHOTS_DIR / request.node.name
    scenario_dir.mkdir(parents=True, exist_ok=True)
    ts = datetime.now().strftime("%Y%m%d_%H%M%S")
    path = str(scenario_dir / f"FAIL_{tag}_{ts}.png")
    try:
        page.screenshot(path=path, full_page=True)
    except Exception:
        pass
    attach_png(path, name=f"失败截图-{tag}")
    # DOM 快照: 定位器失效时直接看到当时的页面结构
    try:
        attach_text(page.content(), name=f"DOM快照-{tag}",
                    attachment_type=allure.attachment_type.HTML)
    except Exception:
        pass
    return path


# ═══════════════════════════════════════════════════════════════════
# 快速关闭 (Windows)
# ═══════════════════════════════════════════════════════════════════

def _driver_proc(p):
    """取 sync_playwright 对象内部的 node 驱动子进程对象。"""
    try:
        return p._impl_obj._connection._transport._proc
    except Exception:
        return None


def kill_browser_tree(driver_pid: int) -> list:
    """杀掉 node 驱动直接拉起的浏览器主进程树（chrome / chrome-headless-shell）。"""
    query = (
        "Get-CimInstance Win32_Process | Where-Object { "
        f"$_.ParentProcessId -eq {driver_pid} -and "
        "($_.Name -like 'chrome*' -or $_.Name -like 'headless*') } "
        "| Select-Object -ExpandProperty ProcessId"
    )
    try:
        ps = subprocess.run(
            ["powershell", "-NoProfile", "-Command", query],
            capture_output=True, text=True, timeout=20,
        )
    except Exception:
        return []
    pids = [x.strip() for x in ps.stdout.split() if x.strip().isdigit()]
    for pid in pids:
        subprocess.run(["taskkill", "/F", "/T", "/PID", pid], capture_output=True, timeout=20)
    return pids


def fast_shutdown(p) -> None:
    """强杀浏览器进程树 + node 驱动，跳过 Playwright 的 30s 优雅关闭等待。

    必须在浏览器已用完、且后续不再使用 p 的操作时调用。
    """
    proc = _driver_proc(p)
    if proc is None:
        return
    try:
        kill_browser_tree(proc.pid)
    finally:
        try:
            proc.kill()
        except Exception:
            pass