基于Python3的接口自动化总结（三）——文件断点续传接口

基于nginx-upload-module模块实现的文件断点续传接口自动化

#! /usr/bin/env python3
# -*-coding:utf-8-*-
import os
import requests
import hashlib


class Upload:
    @staticmethod
    def _get_session_id(file_path):
        """
        获取文件md5值，作为session_id（十六进制）
        :param:  file_path:文件路径（带文件名）
        :return: md5：文件md5值（十六进制）
        """
        assert os.path.exists(file_path), "文件不存在！"
        with open(file_path, 'rb') as f:
            return hashlib.md5(f.read()).hexdigest()

    @staticmethod
    def _get_file_size(file_path):
        """
        获取文件大小
        :param:  file_path:文件路径（带文件名）
        :return: file_size：文件大小
        """
        assert os.path.exists(file_path), "文件不存在！"
        return os.path.getsize(file_path)

    @staticmethod
    def _partition_upload_file_by_rest(file_path, file_size, headers, upload_url, segment_size):
        """
        离线文件分段上传
        :param: file_path:文件路径
        :param: file_size:文件大小
        :param: headers:分段上传信息头
        :param: upload_url:分段上传uri
        :param: segment_size:每个分段大小
        """
        last_size = file_size
        pos = 0
        status_code = 0
        count = 1
        f = open(file_path, 'rb')
        while last_size > 0:
            if last_size >= segment_size:
                x_content_range = f"bytes {pos}-{pos + segment_size - 1}/{file_size}"
                read_size = segment_size
                pos += segment_size
                last_size -= segment_size
            else:
                x_content_range = f"bytes {pos}-{pos + last_size - 1}/{file_size}"
                read_size = last_size
                last_size = 0

            headers['X-Content-Range'] = x_content_range
            response = requests.post(url=upload_url, data=f.read(read_size), headers=headers)
            assert response.status_code in [200, 201], f"第{count}段离线文件上传失败！"  # 接口定义：前几段返回201，最后一段为200
            status_code = response.status_code
            count += 1
        f.close()
        assert status_code is 200, "离线文件上传失败！状态码 != 200"

    def upload_offline_file(self, file_path, upload_url, segment_size):
        """
        断点续传
        :param: file_path:文件路径
        """
        session_id = self._get_session_id(file_path)
        file_size = self._get_file_size(file_path)
        filename = os.path.basename(file_path)
        headers = {
            "Session-ID": session_id,  # 文件id
            "Content-Length": f"{segment_size}",  # 分片长度
            "X-Content-Range": "",  # 文件分片所处范围
            "Content-Type": "application/octet-stream",
            "Content-Disposition": f"attachment;filename={filename}"
        }
        # 分段上传文件
        self._partition_upload_file_by_rest(file_path, file_size, headers, upload_url, segment_size)


if __name__ == "__main__":
    FILE_PATH = r"./test.avi"
    UPLOAD_URL = "http://localhost:80/app/upload"
    SEGMENT_SIZE = 10485760  # 单个片段上传的字节数（默认10M）
    Upload().upload_offline_file(FILE_PATH, UPLOAD_URL, SEGMENT_SIZE)