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

Question

使用后@patch在我的装饰器上它不再起作用了。我想进行一个失败并引发异常的调用，以便我可以检查我的装饰器是否正在捕获此异常，并正在调用某个函数。
mock 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

测试.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()

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