unittest+selenium实现断言失败截图

在用unittest断言的时候，会抛出一个self.failureException异常

 

所以我们只需要重写这个异常类就可以实现在抛出异常之前进行截图：

#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 += '\n失败截图文件: %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()
    
    

断言失败效果：