tw/tests/test_process.py

"""
端到端测试：模拟客户付款后的自动图片处理流程
运行：python test_process.py
"""
import sys
import io
sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding="utf-8", errors="replace")
sys.stderr = io.TextIOWrapper(sys.stderr.buffer, encoding="utf-8", errors="replace")

import asyncio
import os
from dotenv import load_dotenv

load_dotenv()

# ── 测试参数 ────────────────────────────────────────────────
TEST_CUSTOMER_ID = "test_user_001"
TEST_ACC_ID      = "小威哥1216"
TEST_IMAGE_URL   = (
    "https://img.alicdn.com/imgextra/i3/O1CN01tqQst21qIEOdcUOCQ"
    "_!!4611686018427380880-0-amp.jpg"
)
# ────────────────────────────────────────────────────────────


async def step1_analyze():
    """Step 1: 图片分析（模拟 analyze_image 工具调用）"""
    print("\n" + "="*60)
    print("Step 1: 图片分析")
    print("="*60)
    from image.image_analyzer import image_analyzer
    result = await image_analyzer.analyze(TEST_IMAGE_URL)
    print(f"分析结果: {result}")
    return result


async def step2_create_task(analysis: dict):
    """Step 2: 创建 Workflow 任务（模拟 image_analysis_result）"""
    print("\n" + "="*60)
    print("Step 2: 创建 Workflow 任务")
    print("="*60)
    from core.workflow import workflow
    await workflow.image_analysis_result(
        customer_id  = TEST_CUSTOMER_ID,
        image_url    = TEST_IMAGE_URL,
        complexity   = analysis.get("complexity", "normal"),
        acc_id       = TEST_ACC_ID,
        acc_type     = "AliWorkbench",
        gemini_prompt= analysis.get("gemini_prompt", ""),
        aspect_ratio = analysis.get("aspect_ratio", "1:1"),
        perspective  = analysis.get("perspective", "no"),
        proc_type    = analysis.get("proc_type", ""),
        subject      = analysis.get("subject", ""),
        quality      = analysis.get("quality", ""),
    )

    task_id = workflow.customer_active_task.get(TEST_CUSTOMER_ID)
    if task_id:
        task = workflow.tasks[task_id]
        print(f"任务已创建: {task_id[:8]}...")
        print(f"  requirements: {task.requirements}")
        print(f"  image: {task.original_image[:80]}...")
    else:
        print("⚠️  任务创建失败！")
        sys.exit(1)
    return task_id


async def step3_trigger_payment():
    """Step 3: 模拟付款触发处理（等待 Gemini 完成）"""
    print("\n" + "="*60)
    print("Step 3: 模拟付款，触发 Gemini 处理（同步等待完成）")
    print("="*60)
    from core.workflow import workflow

    # 注入发送函数（避免真实发消息，只打印）
    async def fake_send(**kw):
        cid = kw.get("customer_id", kw.get("from_id", "?"))
        content = kw.get("content", kw.get("msg", ""))
        print(f"[FAKE SEND -> {cid}] {str(content)[:120]}")

    workflow._send_message = fake_send

    # 直接调用 _auto_process 同步等待，而非后台 task
    task_id = workflow.customer_active_task.get(TEST_CUSTOMER_ID)
    if not task_id:
        print("  找不到待处理任务！")
        return

    print(f"  开始处理任务: {task_id[:8]}...")
    await workflow._auto_process(task_id, acc_id=TEST_ACC_ID, acc_type="AliWorkbench")


async def step4_check_result():
    """Step 4: 检查结果"""
    print("\n" + "="*60)
    print("Step 4: 检查处理结果")
    print("="*60)
    result_dir = os.getenv("RESULT_IMAGE_DIR", "results")
    if not os.path.exists(result_dir):
        print(f"结果目录不存在: {result_dir}")
        return

    files = sorted(
        [f for f in os.listdir(result_dir) if f.startswith("result_")],
        key=lambda f: os.path.getmtime(os.path.join(result_dir, f)),
        reverse=True,
    )
    if files:
        latest = files[0]
        path   = os.path.join(result_dir, latest)
        size   = os.path.getsize(path)
        print(f"最新结果文件: {latest}")
        print(f"  大小: {size:,} bytes ({size/1024:.1f} KB)")
    else:
        print("结果目录为空，可能处理失败")


async def main():
    print("=" * 60)
    print("  图片处理流程端到端测试")
    print("=" * 60)
    print(f"测试图片: {TEST_IMAGE_URL[:80]}...")

    try:
        analysis = await step1_analyze()
        await step2_create_task(analysis)
        await step3_trigger_payment()
        await step4_check_result()
        print("\n[OK] 测试完成")
    except KeyboardInterrupt:
        print("\n测试被中断")
    except Exception as e:
        import traceback
        print(f"\n[FAIL] 测试失败: {e}")
        traceback.print_exc()


if __name__ == "__main__":
    asyncio.run(main())