我正在编写我的第一个python包,我想为以下装饰器编写单元测试:
class MaxTriesExceededError(Exception): pass def tries(max_tries=3, error_message=os.strerror(errno.ETIME)): def decorator(func): try_count = 0 def wrapper(*args, **kwargs): try_count+=1 try: if try_count <= max_tries: result = func(*args,**kwargs) return result else: raise MaxTriesExceededError(error_message) except: if try_count <= max_tries: wrapper(*args,**kwargs) else: raise Exception return wraps(func)(wrapper) return decorator
装饰器的目的是在函数失败超过max_tries时抛出错误,但是如果没有超过最大尝试次数则吃错误并再试一次.说实话,我不确定代码没有错误.我的问题是双重的,代码是否正确,我如何使用unittest为它编写单元测试?
这是一个更正版本,具有单元测试:
class MaxTriesExceededError(Exception): pass def tries(max_tries=3, error_message="failure"): def decorator(func): def wrapper(*args, **kwargs): for try_count in range(max_tries): try: return func(*args,**kwargs) except: pass raise MaxTriesExceededError(error_message) return wrapper return decorator import unittest class TestDecorator(unittest.TestCase): def setUp(self): self.count = 0 def test_success_single_try(self): @tries(1) def a(): self.count += 1 return "expected_result" self.assertEqual(a(), "expected_result") self.assertEqual(self.count, 1) def test_success_two_tries(self): @tries(2) def a(): self.count += 1 return "expected_result" self.assertEqual(a(), "expected_result") self.assertEqual(self.count, 1) def test_failure_two_tries(self): @tries(2) def a(): self.count += 1 raise Exception() try: a() self.fail() except MaxTriesExceededError: self.assertEqual(self.count,2) def test_success_after_third_try(self): @tries(5) def a(): self.count += 1 if self.count==3: return "expected_result" else: raise Exception() self.assertEqual(a(), "expected_result") self.assertEqual(self.count, 3) if __name__ == '__main__': unittest.main()