具有装饰器的模拟功能。再次使用相同的装饰器来装饰 Mock 对象并使其保持为 Mock

Posted 2023-02-23

【中文标题】具有装饰器的模拟功能。再次使用相同的装饰器来装饰 Mock 对象并使其保持为 Mock

【英文标题】：Mock function that had decorator. Use same decorator again to decorate Mock object and keep it being Mock

【发布时间】：2020-08-25 05:35:50

【问题描述】：

在我的装饰器上使用@patch 后，它不再工作了。我想进行一个失败并引发异常的调用，这样我就可以检查我的装饰器是否捕获了这个异常，并且正在调用某个函数。 模拟 do_sth_in_db 并让它引发异常是很容易的部分，但是在模拟这个方法之后，它不再被修饰 - 所以即使它引发异常，也不会发生任何事情，因为它没有 try/except阻止了。

TLDR：我想在我的模拟函数上放回 @decorator。

my.py

from my_decorator import transaction


class MyClass():

    @transaction
    def do_sth_in_db(self):
        print('Did something in DB')

my_decorator.py

import functools

def rollback_func():
    print('666 rollback fun called')

def push_to_db_func():
    print('777 I have changed database')

def transaction(func):
    functools.wraps(func)
    def wrapper(*args,**kwargs):
        try:
            func(*args, **kwargs)
            push_to_db_func()
            print('worked')
        except Exception:
            rollback_func()
            print('not worked, did rollback')
    return wrapper

test.py

import unittest

from mock import Mock, patch, MagicMock
from my import MyClass
from my_decorator import transaction

class TestMyRollback(unittest.TestCase):

    @patch('my.MyClass.do_sth_in_db')
    @patch('my_decorator.rollback_func')
    @patch('my_decorator.push_to_db_func')
    def test_rollback(self, push_to_db_func_mock, roll_back_func_mock, do_sth_in_db_mock):

        cons = MyClass()

        cons.do_sth_in_db()

        do_sth_in_db_mock.assert_called_once()

        ## needs decorator to work 
        #push_to_db_func_mock.assert_called_once()
        #roll_back_func_mock.assert_not_called()

        ## 
        #roll_back_func_mock.assert_called_once()
        #push_to_db_func_mock.assert_not_called()

【参考方案1】：

How to mock a decorated function 中描述了执行此操作的方法。由于可能不完全清楚如何将其应用于当前问题，这里是工作代码：

import unittest
from unittest.mock import patch

from my_decorator import transaction
from my import MyClass


class TestMyRollback(unittest.TestCase):

    @patch('my.MyClass.do_sth_in_db')
    @patch('my_decorator.rollback_func')
    @patch('my_decorator.push_to_db_func')
    def test_rollback(self, push_to_db_func_mock,
                      roll_back_func_mock, do_sth_in_db_mock):

        # this line re-applies the decorator to the mocked function
        MyClass.do_sth_in_db = transaction(do_sth_in_db_mock)
        cons = MyClass()
        cons.do_sth_in_db()
        do_sth_in_db_mock.assert_called_once()

        push_to_db_func_mock.assert_called_once()
        roll_back_func_mock.assert_not_called()

        # test the exception case 
        push_to_db_func_mock.reset_mock()  # have to reset the mocks
        roll_back_func_mock.reset_mock()

        # this is still the mock for the undecorated function
        do_sth_in_db_mock.side_effect = [Exception]
        cons.do_sth_in_db()
        roll_back_func_mock.assert_called_once()
        push_to_db_func_mock.assert_not_called()