Tool_Agent/tools/local/image_stitcher/tests/test_stitch.py

import sys
import os
import base64
import io

# 解决 Windows 控制台中文输出问题
if sys.platform == "win32":
    sys.stdout.reconfigure(encoding="utf-8")

# 把项目根目录加入 sys.path
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from PIL import Image
from stitch_core import stitch_images

OUTPUT_DIR = os.path.join(os.path.dirname(__file__), "output")
os.makedirs(OUTPUT_DIR, exist_ok=True)


def make_color_image(color: str, width: int = 100, height: int = 80) -> str:
    img = Image.new("RGB", (width, height), color)
    buf = io.BytesIO()
    img.save(buf, format="PNG")
    return base64.b64encode(buf.getvalue()).decode()


def save_result(result: dict, filename: str):
    img_data = base64.b64decode(result["image"])
    path = os.path.join(OUTPUT_DIR, filename)
    with open(path, "wb") as f:
        f.write(img_data)
    print(f"  saved: {path}")


def test_horizontal():
    print("Test: horizontal stitch...")
    red = make_color_image("red", 100, 80)
    blue = make_color_image("blue", 120, 80)
    result = stitch_images([red, blue], direction="horizontal", spacing=0)
    assert result["width"] == 220, f"width should be 220, got {result['width']}"
    assert result["height"] == 80, f"height should be 80, got {result['height']}"
    save_result(result, "horizontal.png")
    print("  [PASS] horizontal")


def test_vertical():
    print("Test: vertical stitch...")
    green = make_color_image("green", 100, 80)
    yellow = make_color_image("yellow", 100, 60)
    result = stitch_images([green, yellow], direction="vertical", spacing=10)
    assert result["width"] == 100, f"width should be 100, got {result['width']}"
    assert result["height"] == 150, f"height should be 150, got {result['height']}"
    save_result(result, "vertical.png")
    print("  [PASS] vertical")


def test_grid():
    print("Test: grid stitch...")
    imgs = [make_color_image(c, 50, 50) for c in ["red", "green", "blue", "yellow"]]
    result = stitch_images(imgs, direction="grid", columns=2, spacing=5)
    expected_w = 50 * 2 + 5 * 1  # 105
    expected_h = 50 * 2 + 5 * 1  # 105
    assert result["width"] == expected_w, f"width should be {expected_w}, got {result['width']}"
    assert result["height"] == expected_h, f"height should be {expected_h}, got {result['height']}"
    save_result(result, "grid.png")
    print("  [PASS] grid")


def test_resize_fit_width():
    print("Test: resize_mode=fit_width...")
    img1 = make_color_image("cyan", 100, 80)
    img2 = make_color_image("magenta", 60, 40)
    result = stitch_images([img1, img2], direction="horizontal", resize_mode="fit_width")
    # fit_width unifies to max width=100, result width = 100+100 = 200
    assert result["width"] == 200, f"width should be 200, got {result['width']}"
    save_result(result, "fit_width.png")
    print("  [PASS] fit_width")


def test_min_images_error():
    print("Test: error on < 2 images...")
    try:
        img = make_color_image("red")
        stitch_images([img], direction="horizontal")
        assert False, "Should raise ValueError"
    except ValueError as e:
        print(f"  [PASS] correct error: {e}")


if __name__ == "__main__":
    print("=== Image Stitcher Tests ===\n")
    test_horizontal()
    test_vertical()
    test_grid()
    test_resize_fit_width()
    test_min_images_error()
    print("\nAll tests passed!")