aiData/Test/TestFirstPage/T_XinDianTu_FirstPageScrollTest.py

import asyncio
import os
import sys
import time
import shutil

import uiautomator2 as u2


project_root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
if project_root not in sys.path:
    sys.path.append(project_root)


from Apps.XinDianTu.Kit import setup_logger, take_screenshot
from Apps.XinDianTu.Config.Setting import SCROLL_DISTANCE_RATIO, WAIT_AFTER_SCROLL
from Apps.XinDianTu.FirstPageKit import run_ocr_rect, run_batch_in_dir


LOGS_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "Logs")
os.makedirs(LOGS_DIR, exist_ok=True)
LOG_FILE_PATH = os.path.join(LOGS_DIR, "T_XinDianTu_FirstPageScrollTest.log")

logger = setup_logger(
    "XinDianTu.FirstPageScrollTest", log_file=LOG_FILE_PATH, clear_old_log=True
)


async def capture_list_page(d, base_dir: str, index: int):
    ts = int(time.time())
    image_uuid = f"xdt_firstpage_{index}_{ts}"
    list_path = take_screenshot(d, image_uuid, save_dir=base_dir)
    logger.info(f"[{index}] 列表截图路径: {list_path}")

    try:
        await run_ocr_rect(list_path, log_path=LOG_FILE_PATH)
        logger.info(f"[{index}] 已根据 OCR+LLM 算法生成首屏绿框调试图。")
    except Exception as e:
        logger.exception(f"[{index}] 运行 OCR 绿框测试失败: {e}")


async def run_scroll_capture():
    logger.info("请先手动打开微信并进入新电途小程序的场站列表第一页。")

    d = u2.connect()
    w, h = d.window_size()
    device_info = d.info or {}
    device_info["width"] = w
    device_info["height"] = h

    logger.info(f"当前设备: {device_info.get('productName')} | 分辨率: {w}x{h}")

    base_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "Output")
    if os.path.exists(base_dir):
        for name in os.listdir(base_dir):
            path = os.path.join(base_dir, name)
            if os.path.isfile(path) or os.path.islink(path):
                os.remove(path)
            elif os.path.isdir(path):
                shutil.rmtree(path)
    os.makedirs(base_dir, exist_ok=True)

    for i in range(1, 6):
        await capture_list_page(d, base_dir, i)
        if i < 5:
            logger.info(f"[{i}] 向上滑动，scale={SCROLL_DISTANCE_RATIO}，等待 {WAIT_AFTER_SCROLL} 秒后继续截图。")
            d.swipe_ext("up", scale=SCROLL_DISTANCE_RATIO)
            await asyncio.sleep(WAIT_AFTER_SCROLL)


async def run_batch_local():
    base_dir = os.path.dirname(os.path.abspath(__file__))
    logger.info(f"批量处理目录中的 JPG 文件: {base_dir}")
    await run_batch_in_dir(base_dir, log_path=LOG_FILE_PATH)


async def run_single_image(image_path: str):
    if not image_path:
        logger.error("单图模式下需要提供图片路径。")
        return
    if not os.path.isabs(image_path):
        image_path = os.path.abspath(image_path)
    if not os.path.exists(image_path):
        logger.error(f"指定的图片不存在: {image_path}")
        return
    logger.info(f"单图 OCR 测试，图片路径: {image_path}")
    await run_ocr_rect(image_path, log_path=LOG_FILE_PATH)


if __name__ == "__main__":
    mode = "scroll"
    image_path = None
    if len(sys.argv) >= 2:
        arg1 = sys.argv[1]
        lower = arg1.lower()
        if lower == "batch":
            mode = "batch"
        elif lower == "image":
            mode = "image"
            if len(sys.argv) >= 3:
                image_path = sys.argv[2]
        else:
            if any(lower.endswith(ext) for ext in [".jpg", ".jpeg", ".png"]):
                mode = "image"
                image_path = arg1
    try:
        if mode == "batch":
            asyncio.run(run_batch_local())
        elif mode == "image":
            asyncio.run(run_single_image(image_path))
        else:
            asyncio.run(run_scroll_capture())
    except KeyboardInterrupt:
        logger.info("用户手动中断测试脚本。")
    except Exception as e:
        logger.exception(f"测试脚本运行异常: {e}")