时间戳

首页 > 解决方案 > 带有自我参数的装饰器的单元测试

python - 带有自我参数的装饰器的单元测试

问题描述

我有一个基于引发异常的重试装饰器：

def retry(times, exceptions):
    def decorator(func):
        def wrapped_func(self):
            attempt = 0
            while attempt < times:
                try:
                    return func(self)
                except exceptions:
                    self.logger.info(
                        "Error, retry in process.. attempt {} of {}".format(
                            attempt, times
                        )
                    )
                    attempt += 1
            return func(self)
        return wrapped_func
    return decorator

并将用于类函数中，例如：

@retry(times=10, exceptions=(ValueError,))
def function_that_fails(self):
    raise ValueError

我正在尝试测试此功能，例如：

class sample_class:
    def __init__(self, logger):
        self._logger = logger
        self._count = 0
    
    @property
    def count(self):
        return self._count

    @property
    def logger(self):
        return self._logger
    
    @retry(times=10, exceptions=(ValueError,))
    def sample_func(self):
        self._count += 1
        raise ValueError

class TestUtils(TestCase):
    def setUp(self):
        self._logger = Mock()
        self._sample_class = sample_class(self._logger)
        self._count = 0
    
    def test_retry_function(self):
        with self.assertRaises(ValueError):
            self._sample_class.sample_func()
        self.assertEqual(self._sample_class.count, 10)

    def test_retry_function_2(self):
        @retry(times=10, exceptions=(ValueError,))
        def sample_func():
            self._count += 1
            raise ValueError
        with self.assertRaises(ValueError):
            sample_func()
        self.assertEqual(self._count, 10)

第一个测试有效，但我想避免使用虚拟类，第二个测试我收到错误：

E   TypeError: wrapped_func() missing 1 required positional argument: 'self'

由于重试函数内部的参数，这是预期的self，我尝试了多个选项但没有成功，我做错了什么？

标签： python

解决方案

您可以简单地使用Mock对象作为self参数，而不是创建整个类：

from unittest.mock import Mock

def test_retry_function_2(self):
    @retry(times=10, exceptions=(ValueError,))
    def sample_func(mock_arg):
        mock_arg._count += 1
        raise ValueError
    mock = Mock(_count=0)
    with self.assertRaises(ValueError):
        sample_func(mock)
    self.assertEqual(mock._count, 10)

更简单的是，您可以使用 Mock 作为修饰函数，并使用内置函数call_count代替_count您必须自己增加的 a：

def test_retry_function_3(self):
    mock = Mock(side_effect=ValueError)
    with self.assertRaises(ValueError):
        retry(10, ValueError)(mock)(Mock())
    self.assertEqual(mock.call_count, 11)

验证您的装饰器是否logger.info在每次重试时调用更容易：

def test_retry_function_4(self):
    mock_func = Mock(side_effect=ValueError)
    mock_self = Mock()
    with self.assertRaises(ValueError):
        retry(10, ValueError)(mock_func)(mock_self)
    self.assertEqual(mock_func.call_count, 11)
    self.assertEqual(mock_self.logger.info.call_count, 10)

推荐阅读