Selenium 自动化下载文件：判断文件下载成功，重试机制，文件重命名

下面是修改后的代码实现：

import os
import re
import time
from selenium.webdriver.common.by import By
from selenium.webdriver.common.action_chains import ActionChains

def click_download_files(self, urn, el):
    '文件点击下载'
    file_names = []
    download_links = self.driver.find_elements(By.CSS_SELECTOR, el)
    try:
        for link in download_links:
            file_title = link.get_attribute('title')
            t_f_n = re.findall(REG, file_title)

            if len(t_f_n) == 0:
                QueryResult.update(UPDATE_INVOICES_SQL, (generate_utils.get_strtime(), urn))
                continue
            
            # 2023-3-27 测试中发现存在文件名为空的错误数据导致程序死循环
            double_check = r'^.'
            tfn = re.findall(double_check, t_f_n[0].strip().strip())
            if len(tfn) > 0:
                QueryResult.update(UPDATE_INVOICES_SQL, (generate_utils.get_strtime(), urn))
                continue
            
            file_title = t_f_n[0].strip()
            new_file_title = urn + '__' + t_f_n[0].strip()  # 修改后的文件名

            if link.is_enabled():  # 文件是否可点击下载
                ActionChains(self.driver).move_to_element(link).click().perform()

                # 文件下载时间
                err_flag, seconds = self.download_wait(self.get_download_path(), file_title, 600)
                if not err_flag:
                    logger.info('
 >> 【urn = {}】附件下载完成，文件名: 【{}】，用时: 【{}】'.format(urn, file_title, seconds))
                    os.rename(self.get_download_path() + '\' + file_title,
                              self.get_download_path() + '\' + new_file_title)
                    file_names.append(new_file_title)
                    logger.info('
 >> 【urn = {}】防止冲突，文件改名: old = 【{}】, new = 【{}】'.format(urn, file_title, new_file_title))
                else:
                    # 下载超时
                    ValidateData().add_record(urn, self.name, Operate.DOWNLOAD.value,
                                              file_title + ' ' + ErrorMsg.DOWNLOAD_FILE_TIMEOUT, el)
                    # ValidateData().update_urn_is_load(urn)
                    QueryResult.update(UPDATE_INVOICES_1_SQL, (generate_utils.get_strtime(), urn))
            else:
                # 文件不可点击下载
                ValidateData().add_record(urn, self.name, Operate.DOWNLOAD.value,
                                          file_title + ' ' + ErrorMsg.THE_DOWNLOAD_FILE_DOES_NOT_EXIST, el)
                ValidateData().update_urn_is_load(urn)
    
    # 重试下载失败的文件
    retry_count = 0
    while retry_count < 3:
        failed_files = []
        for file_name in file_names:
            if not os.path.exists(self.get_download_path() + '\' + file_name):
                failed_files.append(file_name)
        
        if len(failed_files) == 0:
            break
        
        for failed_file in failed_files:
            link = self.driver.find_element(By.CSS_SELECTOR, el)
            file_title = link.get_attribute('title')
            t_f_n = re.findall(REG, file_title)
            
            if len(t_f_n) == 0:
                QueryResult.update(UPDATE_INVOICES_SQL, (generate_utils.get_strtime(), urn))
                continue
            
            double_check = r'^.'
            tfn = re.findall(double_check, t_f_n[0].strip().strip())
            if len(tfn) > 0:
                QueryResult.update(UPDATE_INVOICES_SQL, (generate_utils.get_strtime(), urn))
                continue
            
            file_title = t_f_n[0].strip()
            new_file_title = urn + '__' + t_f_n[0].strip()  # 修改后的文件名
            
            if link.is_enabled():  # 文件是否可点击下载
                ActionChains(self.driver).move_to_element(link).click().perform()
                
                # 文件下载时间
                err_flag, seconds = self.download_wait(self.get_download_path(), file_title, 600)
                if not err_flag:
                    logger.info('
 >> 【urn = {}】附件下载完成，文件名: 【{}】，用时: 【{}】'.format(urn, file_title, seconds))
                    os.rename(self.get_download_path() + '\' + file_title,
                              self.get_download_path() + '\' + new_file_title)
                    file_names.append(new_file_title)
                    logger.info('
 >> 【urn = {}】防止冲突，文件改名: old = 【{}】, new = 【{}】'.format(urn, file_title, new_file_title))
                else:
                    # 下载超时
                    ValidateData().add_record(urn, self.name, Operate.DOWNLOAD.value,
                                              file_title + ' ' + ErrorMsg.DOWNLOAD_FILE_TIMEOUT, el)
                    # ValidateData().update_urn_is_load(urn)
                    QueryResult.update(UPDATE_INVOICES_1_SQL, (generate_utils.get_strtime(), urn))
            else:
                # 文件不可点击下载
                ValidateData().add_record(urn, self.name, Operate.DOWNLOAD.value,
                                          file_title + ' ' + ErrorMsg.THE_DOWNLOAD_FILE_DOES_NOT_EXIST, el)
                ValidateData().update_urn_is_load(urn)
        
        retry_count += 1
    
    # 对每个文件修改名字
    for file_name in file_names:
        os.rename(self.get_download_path() + '\' + file_name,
                  self.get_download_path() + '\' + file_name + '_new')
    
    # 判断文件是否下载完成
    for file_name in file_names:
        if os.path.exists(self.get_download_path() + '\' + file_name + '_new'):
            logger.info('
 >> 【urn = {}】文件下载完成，文件名: 【{}】'.format(urn, file_name))
        else:
            logger.info('
 >> 【urn = {}】文件下载失败，文件名: 【{}】'.format(urn, file_name))


def download_wait(self, directory, temp_file_name, max_wait_time_sec):
    '''
    等待指定文件出现并记录等待的时间，最多等待 10 分钟。如果文件在 10 分钟内出现，则返回 err_flag 为 False，
    否则返回 err_flag 为 True。
    directory 文件路径
    temp_file_name 文件名
    '''
    seconds = 0
    dl_wait = True
    err_flag = False

    while dl_wait:
        time.sleep(5)
        files = os.listdir(directory)

        for nm in files:
            if nm == temp_file_name:
                dl_wait = False

        seconds += 5

        if seconds > max_wait_time_sec:
            dl_wait = False
            err_flag = True

    return err_flag, seconds