RagFlow_Plugs/test/test_image_compression_fix.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
测试图片压缩修复
验证修改后的压缩方法是否能成功将图片压缩到5000KB以内
"""
import sys
import os
from io import BytesIO
from PIL import Image
import random

# 添加项目根目录到Python路径
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from utils.minio.image_util import ImageUtil

def create_large_test_image(width=3000, height=3000) -> BytesIO:
    """
    创建一个大尺寸测试图片
    
    Args:
        width: 图片宽度
        height: 图片高度
        
    Returns:
        BytesIO: 大尺寸图片流
    """
    print(f"创建 {width}x{height} 的测试图片...")
    
    # 创建一个大尺寸图片，使用随机颜色填充
    img = Image.new('RGB', (width, height), color=(random.randint(0, 255), random.randint(0, 255), random.randint(0, 255)))
    
    # 将图片保存到BytesIO流
    img_stream = BytesIO()
    img.save(img_stream, format='PNG')
    img_stream.seek(0)
    
    # 检查图片大小
    img_stream.seek(0, 2)
    size_kb = img_stream.tell() / 1024
    img_stream.seek(0)
    
    print(f"测试图片创建完成，大小为 {size_kb:.2f}KB")
    return img_stream

def test_image_compression():
    """
    测试图片压缩方法
    """
    print("开始测试图片压缩方法...")
    
    # 创建ImageUtil实例
    image_util = ImageUtil()
    
    # 测试不同尺寸的图片压缩
    test_sizes = [
        (3000, 3000),   # 约 25MB
        (4000, 4000),   # 约 45MB
        (5000, 5000)    # 约 70MB
    ]
    
    for width, height in test_sizes:
        print(f"\n=== 测试 {width}x{height} 图片压缩 ===")
        
        # 创建大尺寸测试图片
        img_stream = create_large_test_image(width, height)
        
        # 调用压缩方法
        compressed_stream = image_util._compress_image(img_stream, "test_large_image.png")
        
        # 检查压缩后的大小
        compressed_stream.seek(0, 2)
        compressed_size_kb = compressed_stream.tell() / 1024
        compressed_stream.seek(0)
        
        print(f"压缩后大小: {compressed_size_kb:.2f}KB")
        
        # 验证压缩结果
        if compressed_size_kb <= 5000:
            print("✅ 压缩成功！压缩后大小小于等于5000KB")
        else:
            print("❌ 压缩失败！压缩后大小仍大于5000KB")
    
    # 测试_compress_image_to_bytes方法
    print(f"\n=== 测试 _compress_image_to_bytes 方法 ===")
    img_stream = create_large_test_image(4000, 4000)
    compressed_bytes = image_util._compress_image_to_bytes(img_stream)
    compressed_size_kb = len(compressed_bytes) / 1024
    print(f"压缩后字节大小: {compressed_size_kb:.2f}KB")
    
    if compressed_size_kb <= 5000:
        print("✅ _compress_image_to_bytes 压缩成功！")
    else:
        print("❌ _compress_image_to_bytes 压缩失败！")
    
    print("\n=== 所有测试完成 ===")

if __name__ == "__main__":
    test_image_compression()