python - 如何在 python 包中测试装饰器

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