Mock function that had decorator. Use same decorator again to decorate Mock object and keep it being Mock

Question

After using @patch on my decorator it does not work anymore. I want to make a call that will fail and raise an exception so I can check if my decorator is catching this exception, and is calling some function.
Mocking do_sth_in_db and getting it to raise an exception is the easy part, but after mocking this method, it is not decorated any more - so even if it raises an exception, nothing will happen because it doesn't have the try/except block anymore.

TLDR: I want to put back @decorator on my mocked function.

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

score 2 · Answer 1 · answered May 10 '20 at 12:33

The way to do this is described in How to mock a decorated function. As it may not be completely clear how to apply this to the current problem, here is the working code:

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

Mock function that had decorator. Use same decorator again to decorate Mock object and keep it being Mock

1 Answers1