Python mock: replace function with another function if sys.argv argument is set

Question

I have a snippet of code that looks like this:

if args.module == 'omega':
    with mock.patch.object(quantum_entangler, 'sort_by_dirac', sort_by_heisenberg):
        quantum_entangler.main(atom_no)
else:
    quantum_entangler.main(atom_no)

Which basically lets me set sort function with --module option in the command line call.

I was wondering if there's a way to implement this with just the mock.patch.object() line without the if/else block.

Answer 1

You can use a dictionary of references to the sorting function

The idea is to create a dictionary which contains references to the sorting functions.

The following code is found in a file called test.py. To execute test.py I have to pass it 2 arguments by command line:

• --module with one of the value omega | dirac
  NOTE. I have defined the value dirac to execute the normal sorting by the function sort_by_dirac(), while by omega is selected the sorting by the function sort_by_heisenberg() as you already do in your code
• --atom_no which can be of any value (it is used to pass something to main() function of the module quantum_entangler.py):

Code of test.py:

from unittest import mock
import argparse
import quantum_entangler
from quantum_entangler import sort_by_heisenberg, sort_by_dirac

parser = argparse.ArgumentParser()
parser.add_argument("--module")
parser.add_argument("--atom_no")
args = parser.parse_args()

atom_no = args.atom_no

# THIS IS THE DICTIONARY of references
sortFunction = {'omega': sort_by_heisenberg, 'delta': sort_by_dirac}

# ---> here you find patch.object() function WITHOUT if/else as you request
# in the question
with mock.patch.object(quantum_entangler,'sort_by_dirac',sortFunction[args.module]):
    quantum_entangler.main(atom_no)

Execution of test.py

To test the code I have written a module quantum_entangler.py with the following 3 functions:

def sort_by_dirac(atom_no):
    print(f"execute sort_by_dirac({atom_no})")
    
def sort_by_heisenberg(atom_no):
    print(f"execute sort_by_heisenberg({atom_no})")
    
def main(atom_no):
    sort_by_dirac(atom_no)

Execution of test.py with --module dirac:

python test.py --atom_no 10 --module dirac

# output
execute sort_by_dirac(10)

Execution of test.py with --module omega:

python test.py --atom_no 10 --module omega

# output
execute sort_by_heisenberg(10)