RagFlow_Plugs/test/test_image_compression_bytes.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_to_bytes():
    """
    测试图片压缩到字节流功能
    """
    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_bytes = image_util._compress_image_to_bytes(img_stream, "test_image.png", max_size_kb=5000)
    
    # 检查压缩后大小
    compressed_size = len(compressed_bytes) / 1024
    print(f"压缩后大小: {compressed_size:.2f}KB")
    
    # 验证压缩后大小
    assert compressed_size <= 5000, f"压缩后大小 {compressed_size:.2f}KB 超过了最大限制 5000KB"
    
    # 验证返回类型
    assert isinstance(compressed_bytes, bytes), f"返回类型应为bytes，实际为 {type(compressed_bytes)}"
    
    print("图片压缩到字节流测试成功!")

if __name__ == "__main__":
    test_image_compression_to_bytes()