文章是参考了简书文章,原文有更详细和全面的方法,在此是搬运了我用的两种,原文链接:https://www.jianshu.com/p/2c471acfea2e
1 # coding=utf-8 2 import sys 3 import functools 4 import traceback 5 import inspect 6 import unittest 7 8 9 def retry(target=None, max_n=1, func_prefix="test"): 10 """ 11 一个装饰器,用于unittest执行测试用例出现失败后,自动重试执行 12 13 # example_1: test_001默认重试1次 14 class ClassA(unittest.TestCase): 15 @retry 16 def test_001(self): 17 raise AttributeError 18 19 20 # example_2: max_n=2,test_001重试2次 21 class ClassB(unittest.TestCase): 22 @retry(max_n=2) 23 def test_001(self): 24 raise AttributeError 25 26 27 # example_3: test_001重试3次; test_002重试3次 28 @retry(max_n=3) 29 class ClassC(unittest.TestCase): 30 def test_001(self): 31 raise AttributeError 32 33 def test_002(self): 34 raise AttributeError 35 36 37 # example_4: test_102重试2次, test_001不参与重试机制 38 @retry(max_n=2, func_prefix="test_1") 39 class ClassD(unittest.TestCase): 40 def test_001(self): 41 raise AttributeError 42 43 def test_102(self): 44 raise AttributeError 45 46 47 :param target: 被装饰的对象,可以是class, function 48 :param max_n: 重试次数,没有包含必须有的第一次执行 49 :param func_prefix: 当装饰class时,可以用于标记哪些测试方法会被自动装饰 50 :return: wrapped class 或 wrapped function 51 """ 52 53 def decorator(func_or_cls): 54 if inspect.isfunction(func_or_cls): 55 @functools.wraps(func_or_cls) 56 def wrapper(*args, **kwargs): 57 n = 0 58 while n <= max_n: 59 try: 60 n += 1 61 func_or_cls(*args, **kwargs) 62 return 63 except Exception: # 可以修改要捕获的异常类型 64 if n <= max_n: 65 trace = sys.exc_info() 66 traceback_info = str() 67 for trace_line in traceback.format_exception(trace[0], trace[1], trace[2], 3): 68 traceback_info += trace_line 69 print(traceback_info) # 输出组装的错误信息 70 args[0].tearDown() 71 args[0].setUp() 72 else: 73 raise 74 75 return wrapper 76 elif inspect.isclass(func_or_cls): 77 for name, func in list(func_or_cls.__dict__.items()): 78 if inspect.isfunction(func) and name.startswith(func_prefix): 79 setattr(func_or_cls, name, decorator(func)) 80 return func_or_cls 81 else: 82 raise AttributeError 83 84 if target: 85 return decorator(target) 86 else: 87 return decorator 88 89 90 class Retry(object): 91 """ 92 类装饰器, 功能与Retry一样 93 94 95 # example_1: test_001默认重试1次 96 class ClassA(unittest.TestCase): 97 @Retry 98 def test_001(self): 99 raise AttributeError 100 101 102 # example_2: max_n=2,test_001重试2次 103 class ClassB(unittest.TestCase): 104 @Retry(max_n=2) 105 def test_001(self): 106 raise AttributeError 107 108 109 # example_3: test_001重试3次; test_002重试3次 110 @Retry(max_n=3) 111 class ClassC(unittest.TestCase): 112 def test_001(self): 113 raise AttributeError 114 115 def test_002(self): 116 raise AttributeError 117 118 119 # example_4: test_102重试2次, test_001不参与重试机制 120 @Retry(max_n=2, func_prefix="test_1") 121 class ClassD(unittest.TestCase): 122 def test_001(self): 123 raise AttributeError 124 125 def test_102(self): 126 raise AttributeError 127 128 """ 129 130 def __new__(cls, func_or_cls=None, max_n=1, func_prefix="test"): 131 self = object.__new__(cls) 132 if func_or_cls: 133 self.__init__(func_or_cls, max_n, func_prefix) 134 return self(func_or_cls) 135 else: 136 return self 137 138 def __init__(self, func_or_cls=None, max_n=1, func_prefix="test"): 139 self._prefix = func_prefix 140 self._max_n = max_n 141 142 def __call__(self, func_or_cls=None): 143 if inspect.isfunction(func_or_cls): 144 @functools.wraps(func_or_cls) 145 def wrapper(*args, **kwargs): 146 n = 0 147 while n <= self._max_n: 148 try: 149 n += 1 150 func_or_cls(*args, **kwargs) 151 return 152 except Exception: # 可以修改要捕获的异常类型 153 if n <= self._max_n: 154 trace = sys.exc_info() 155 traceback_info = str() 156 for trace_line in traceback.format_exception(trace[0], trace[1], trace[2], 3): 157 traceback_info += trace_line 158 print(traceback_info) # 输出组装的错误信息 159 args[0].tearDown() 160 args[0].setUp() 161 else: 162 raise 163 164 return wrapper 165 elif inspect.isclass(func_or_cls): 166 for name, func in list(func_or_cls.__dict__.items()): 167 if inspect.isfunction(func) and name.startswith(self._prefix): 168 setattr(func_or_cls, name, self(func)) 169 return func_or_cls 170 else: 171 raise AttributeError