Макетная функция с декоратором. Снова используйте тот же декоратор, чтобы украсить объект Mock и оставить его Mock

Question

После использования @patch в моем декораторе он больше не работает. Я хочу сделать вызов, который завершится ошибкой и вызовет исключение, чтобы я мог проверить, перехватывает ли мой декоратор это исключение и вызывает ли некоторую функцию. но после имитации этого метода он больше не украшается - поэтому даже если он вызовет исключение, ничего не произойдет, потому что у него больше нет блока try/except.

TL; DR: я хочу вернуть @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()

Answer 1

Способ сделать это описан в Как имитировать декорированную функцию . Поскольку может быть не совсем понятно, как применить это к текущей проблеме, вот рабочий код:

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