PyTest: test which function was called in the if statement

Question

I have a function, which considers three different cases and for each case calls a different function, like in the example below

def my_fun(input):
 if input == 1:
    fun1()
 if input == 2:
    fun2()
 if input == 3:
    fun3()

I would like to write a test for function my_fun using py.test, but I don't know how I can test if the proper function was called for a given input?

Answer 1

As @fmarc commented, which function is called is less important than testing that my_fun does the right thing. But, you can mock each of the three functions, then test that the correct one was called independently of what each function actually does. (Note that mock is a 3rd-party module in Python 2 that you'll need to install; it is available the standard library in Python 3.3(?) as unittest.mock.)

One simple example:

import mock

def test_my_fun():
    with mock.patch('fun1', wraps=fun1) as mock_fun1:
        with mock.patch('fun2', wraps=fun2) as mock_fun2:
            with mock.patch('fun3', wraps=fun3) as mock_fun3:
                my_fun(1)
                mock_fun1.assert_called_with()
                mock_fun2.assert_not_called()
                mock_fun3.assert_not_called()

Check the documentation of your installation of mock to see which methods are supported.

Answer 2

If all fun1(), fun2(), and fun3() return results I'd go with parametrization.

# Here an example of functions to be tested

def fun1():
    return 1
def fun2():
    return 2
def fun3():
    return 3

def my_fun(input):
    if input == 1:
        return fun1()
    if input == 2:
        return fun2()
    if input == 3:
        return fun3()

Here in the test matrix. The expected result is also a call to a function under test so you can match the function call in the test (as you asked in the question):

# Parametrized test
import pytest

@pytest.mark.parametrize("test_input, expected", [(1, fun1()), (2, fun2()), (3, fun3())])
# ---------------------------------------
# test_input |   1    |   2    |   3    |
# ---------------------------------------
#   expected | fun1() | fun2() | fun3() |
# ---------------------------------------
def test_my_fun(test_input, expected):
    assert my_fun(test_input) == expected

Test run:

mac: py.test -v test_parametrization.py
=================================================================================== test session starts ====================================================================================
platform darwin -- Python 2.7.11, pytest-2.9.1, py-1.4.31, pluggy-0.3.1 -- /Library/Frameworks/Python.framework/Versions/2.7/Resources/Python.app/Contents/MacOS/Python
cachedir: .cache
rootdir: /Users/user/Repos/playground, inifile:
plugins: hidecaptured-0.1.2, instafail-0.3.0
collected 3 items

test_parametrization.py::test_my_fun[1-1] PASSED
test_parametrization.py::test_my_fun[2-2] PASSED
test_parametrization.py::test_my_fun[3-3] PASSED

================================================================================= 3 passed in 0.01 seconds =================================================================================