Unittest .called doesn't work when the method is called through a variable that's referencing it

user_4900457 0
Viewed 48

I have something like this in my code:

def submethod_a():
    pass

def submethod_b():
    pass

_method_registry = {
    "a": submethod_a,
    "b": submethod_b,
}

def main_method(key):
    method_to_use = _method_registry[key]
    method_to_use()

And then I have something like this in my test:

from unittest.mock import patch


@patch("submethod_b")
def test_correct_method_and_arguments(mocked_submethod_b):
    main_method(key)

    assert mocked_submethod_b.called

And the test fails. I did this to test:

def main_method(key):
    if key == "a":
        submethod_a()
    else:
        submethod_b()

And then the test works.

So, .called doesn't work when the method is in a reference? Is this the expected behaviour? What can I do? If I don't solve it any other way, I think I'll refactor the code and use classes with 2 different implementations of "submethod", I think that should work.

Thanks.

python-unittest python-3.x python
1 Answers

When the source code was initially read, it already stored the original submethod_b function within the variable _method_registry, thus by the time you patched it at runtime, although the function was replaced, it didn't change the already-stored value in this variable. To prove it, let's add prints.

main.py

def submethod_a():
    pass

def submethod_b():
    pass

_method_registry = {
    "a": submethod_a,
    "b": submethod_b,
}


def main_method(key):
    print(_method_registry["b"], "- Stored function")
    print(submethod_b, "- Patched function")

test_main.py

from unittest.mock import patch

from main import main_method


def test_no_patch():
    main_method("b")


@patch("main.submethod_b")
def test_with_patch(mocked_submethod_b):
    main_method("b")

Output

$ pytest -q -rP
================================================================================================= PASSES ==================================================================================================
______________________________________________________________________________________________ test_no_patch ______________________________________________________________________________________________
------------------------------------------------------------------------------------------ Captured stdout call -------------------------------------------------------------------------------------------
<function submethod_b at 0x7fba0fb63ee0> - Stored function
<function submethod_b at 0x7fba0fb63ee0> - Patched function
_____________________________________________________________________________________________ test_with_patch _____________________________________________________________________________________________
------------------------------------------------------------------------------------------ Captured stdout call -------------------------------------------------------------------------------------------
<function submethod_b at 0x7fba0fb63ee0> - Stored function
<MagicMock name='submethod_b' id='140437103934432'> - Patched function
2 passed in 0.04s

As you can see, without the patch, the value is just the same original function. But with the patch, notice that only the actual function was replaced, not the one already stored.

Alternative Solution

  1. Put the submethod functions to its own file.
  2. Import the submethods to your main file and use as originally done.
  3. Patch the submethod file.
  4. Reload the main file so that the patched submethods would be the one used when recreating the _method_registry.

main.py

from submethod import submethod_a, submethod_b  # Or: import submethod

_method_registry = {
    "a": submethod_a,  # Depending on chosen import style, this can be: submethod.submethod_a
    "b": submethod_b,  # Depending on chosen import style, this can be: submethod.submethod_b
}


def main_method(key):
    print(_method_registry["b"], "- Stored function")
    print(submethod_b, "- Patched function")

    method_to_use = _method_registry[key]
    method_to_use()

submethod.py

def submethod_a():
    pass

def submethod_b():
    pass

test_main.py

from importlib import reload
import sys

from unittest.mock import patch

from main import main_method


@patch("submethod.submethod_b")
def test_with_patch(mocked_submethod_b):
    # Now that the patch is  in effect and replaced the submethod.submethod_b, reload the main file so that it recreates _method_registry with the patched function
    reload(sys.modules['main'])

    main_method("b")

    assert mocked_submethod_b.called

Output

$ pytest -q -rP
================================================================================================= PASSES ==================================================================================================
_____________________________________________________________________________________________ test_with_patch _____________________________________________________________________________________________
------------------------------------------------------------------------------------------ Captured stdout call -------------------------------------------------------------------------------------------
<MagicMock name='submethod_b' id='140432453077360'> - Stored function
<MagicMock name='submethod_b' id='140432453077360'> - Patched function
1 passed in 0.04s

Now, the stored function and the patched function itself both points to the mocked version. Thus, the .called assertion is now successful.

user_110438250
Related