RagFlow_Plugs/test/test_image_compression.py

#!/usr/bin/env python3
"""
测试图片压缩功能
"""

import os
import sys
from io import BytesIO
from PIL import Image

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

from utils.minio.image_util import image_util

# 生成一个大的测试图片
def generate_test_image(width=2000, height=2000, color=(255, 0, 0)):
    """
    生成一个大的测试图片
    """
    print(f"生成测试图片，大小: {width}x{height}")
    img = Image.new('RGB', (width, height), color=color)
    img_stream = BytesIO()
    img.save(img_stream, format='PNG')
    img_stream.seek(0)
    return img_stream

# 测试图片压缩功能
def test_image_compression():
    """
    测试图片压缩功能
    """
    print("开始测试图片压缩功能...")
    
    # 生成测试图片
    img_stream = generate_test_image()
    
    # 检查压缩前大小
    img_stream.seek(0, 2)
    original_size = img_stream.tell() / 1024
    img_stream.seek(0)
    print(f"压缩前大小: {original_size:.2f}KB")
    
    # 调用压缩方法
    compressed_stream = image_util._compress_image(img_stream, "test_image.png", max_size_kb=5000)
    
    # 检查压缩后大小
    compressed_stream.seek(0, 2)
    compressed_size = compressed_stream.tell() / 1024
    compressed_stream.seek(0)
    print(f"压缩后大小: {compressed_size:.2f}KB")
    
    # 验证压缩后大小
    assert compressed_size <= 5000, f"压缩后大小 {compressed_size:.2f}KB 超过了最大限制 5000KB"
    print("图片压缩测试成功!")

if __name__ == "__main__":
    test_image_compression()