How to mock an object that is an input to a function?

Question

Usually when I mock, I have the following type of setup

# my_script.py
import numpy as np

def my_func(x):
    out = np.power(x, 2)
    return out

then to test the numpy power call in my_script:

# test_myscript.py

import numpy as np
import unittest
import mock

from my_script import my_func


class TestMyScript(unittest.TestCase):

    @mock.patch("my_script.np")
    def test_my_func(self, mock_os):
        """Test that numpy.power was called"""
        a = np.array([1, 2, 3])
        my_func(a)
        mock_os.power.assert_called_with(a, 2)


if __name__ == '__main__':
    unittest.main()

This works fine.

But now if the situation changes, and say I give the numpy module as an argument into my_func; I don't know how to mock numpy in this case.

How would I mock numpy in the function below in the same way as it was mocked in test_myscript above?

Note that numpy will not be imported in my_script.py but will instead be imported in a separate script that runs functions from my_script.py.

# my_script.py
# numpy NOT imported in this script!

def my_func(x, numpy):
    out = numpy.power(x, 2)
    return out

EDIT:

Based on @Daniel Roseman's comment, I am including some more code to be explicit on how the functions are called

# main_script.py

import numpy as np

from my_script import my_func

def main():
    a = np.array([1, 2, 3])
    my_func(a, np)  # numpy is passed into `my_func`

Then to test, I am trying the below

# test_myscript.py

import numpy as np
import unittest
import mock

from my_script import my_func

class TestMyScript(unittest.TestCase):

    @mock.patch("main_script.np")  # import from main_script since numpy is imported here
    def test_my_func(self, mock_os):
        """Test that numpy.power was called"""
        a = np.array([1, 2, 3])
        my_func(a)
        mock_os.power.assert_called_with(a, 2)


if __name__ == '__main__':
    unittest.main()

But this fails with

Ran 1 test in 0.154s
>>> FAILED (failures=1)
>>> AssertionError: Expected 'power' to have been called.

user_4139143 · Answer

I found that using the unittest.mock.Mock object here worked best.

So if we have:

# my_script.py

def my_func(x, numpy):
    out = numpy.power(x, 2)
    return out

Then to test it we have:

# test_myscript.py

import numpy as np
import unittest
from unittest.mock import Mock

from my_script import my_func

numpy_mock = Mock()


class TestMyScript(unittest.TestCase):

    def test_my_func(self):
        """Test that numpy.power was called"""
        a = np.array([1, 2, 3])
        _ = my_func(a, numpy_mock)  # pass the mocked object here
        numpy_mock.power.assert_called_once_with(a, 2)


if __name__ == '__main__':
    unittest.main()

Which passes the test

How to mock an object that is an input to a function?

1 Answers