Question

我正在编写我的第一个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为它编写单元测试？

Answer 1

这是一个更正版本，包含单元测试：

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()

如何测试python包中的装饰器

1 个答案: