unittest实现用例运行失败截图

#coding: utf-8
 
import unittest, random, os, traceback
from selenium import webdriver
 
SCREENSHOT_DIR = 'D:\'
class Test1(unittest.TestCase):
    
    def setUp(self):
        self.driver = webdriver.Chrome('C:\Program Files (x86)\Google\Chrome\Application\chromedriver.exe')
        
        #重新赋值failureException，注意：failureException的值是一个类，不是类实例
        self.failureException = self.failure_monitor()
    
    def failure_monitor(self):
        test_case = self #将self赋值给test_case，以便下方的AssertionErrorPlus内部类可调用外部类的方法
        
        class AssertionErrorPlus(AssertionError):
            def __init__(self, msg):
                try:
                    cur_method = test_case._testMethodName #当前test函数的名称
                    unique_code =  ''.join(random.sample('1234567890',5)) #随机生成一个值，以便区分同一个test函数内不同的截图
                    file_name = '%s_%s.png' % (cur_method, unique_code)
                    test_case.driver.get_screenshot_as_file(os.path.join(SCREENSHOT_DIR, file_name)) #截图生成png文件
                    print('失败截图已保存到: %s' % file_name)
                    msg += '
失败截图文件: %s' % file_name
                except BaseException:
                    print('截图失败: %s' % traceback.format_exc())
                
                super(AssertionErrorPlus, self).__init__(msg)
                
        return AssertionErrorPlus #返回AssertionErrorPlus类
    
    def test1(self):
        self.assertEqual(0, 1, '错误提示')
        
if __name__ == "__main__":
    unittest.main()