import os
import time

def write_large_file(file_path, size_mb):
    """写入大文件"""
    with open(file_path, 'wb') as f:
        f.write(os.urandom(size_mb * 1024 * 1024))  # 写入指定大小的随机数据

def read_large_file(file_path):
    """读取大文件"""
    with open(file_path, 'rb') as f:
        f.read()  # 读取文件内容

def write_small_files(dir_path, num_files, file_size):
    """写入小文件"""
    for i in range(num_files):
        file_path = os.path.join(dir_path, f'small_file_{i}.txt')
        with open(file_path, 'wb') as f:
            f.write(os.urandom(file_size))  # 写入指定大小的随机数据

def read_small_files(dir_path, num_files):
    """读取小文件"""
    for i in range(num_files):
        file_path = os.path.join(dir_path, f'small_file_{i}.txt')
        with open(file_path, 'rb') as f:
            f.read()  # 读取文件内容

def measure_io(test_type, *args):
    """测量IO性能"""
    start_time = time.time()
    if test_type == 'large_write':
        write_large_file(*args)
    elif test_type == 'large_read':
        read_large_file(*args)
    elif test_type == 'small_write':
        write_small_files(*args)
    elif test_type == 'small_read':
        read_small_files(*args)
    end_time = time.time()
    return end_time - start_time

def main():
    dir_path = '/var/lib/docker/'  # 修改为你的测试目录
    large_file_size_mb = 10240  # 大文件大小（MB）
    small_file_size = 1024  # 小文件大小（字节）
    num_small_files = 100  # 小文件数量

    # 大文件测试
    print("开始大文件写入测试...")
    large_write_time = measure_io('large_write', os.path.join(dir_path, 'large_file.bin'), large_file_size_mb)
    print(f"大文件写入时间: {large_write_time:.2f}秒")
    
    print("开始大文件读取测试...")
    large_read_time = measure_io('large_read', os.path.join(dir_path, 'large_file.bin'))
    print(f"大文件读取时间: {large_read_time:.2f}秒")

    # 小文件测试
    print("开始小文件写入测试...")
    small_write_time = measure_io('small_write', dir_path, num_small_files, small_file_size)
    print(f"小文件写入时间: {small_write_time:.2f}秒")

    print("开始小文件读取测试...")
    small_read_time = measure_io('small_read', dir_path, num_small_files)
    print(f"小文件读取时间: {small_read_time:.2f}秒")

if __name__ == '__main__':
    main()