Python unit test: how to test function that calls other function?

Question

I am new to python. I want to test a function that takes raw inputs and calls another function base on the input. I know this can be tested using mock, but I don't know how.

My function is like this:

def show_main_menu(self):
    """shows main menu to user"""
    print '1:create account \t 2:Log in \t any other to exit \n'
    while True:
        try:
            option = int(raw_input('Enter your option: \t'))
        except ValueError:
            print 'please enter valid input \n'
            continue
        else:
            break
    if option == 1:
        self.add_account()
        print '\n ***successfully created account***\n'
        self.menu_post_transaction()
    elif option == 2:
        self.login()
    else:
        print 'Thank you for using our services'
        exit()

How can I test this function using unittest ?

Answer 1

main menu module:

import unittest
from your_tests_directory import add_account_test, login_test

def menu():
    loader = unittest.TestLoader()
    suite = unittest.TestSuite()

    option = int(raw_input('Enter your option: \t'))
    if option == 1:
       module = add_account_test
    else:
        module = login_test

    suite.addTest(loader.loadTestsFromModule(module))
    unittest.TextTestRunner(verbosity=2).run(suite)

if __name__ == '__main__':
    menu()

add_account_test module:

import unittest
from other_module import add_account, menu_post_transaction

class AddAccountTest(unittest.TestCase):
    def test_account(self):
        add_account()
        print '\n ***successfully created account***\n'
        menu_post_transaction()

login_test module:

import unittest
from other_module import login

class LoginTest(unittest.TestCase):
    def test_login(self):
        login()

in 'other_module' you will put all of the actual functions...

Answer 2

I would strongly suggest you refactor your code a bit to make it testable, by just extracting the raw_input related code to a function, you wont need to mock anything, and _get_raw_input() will be reusable.

While testing you just have to pass your own function in the get_option argument which simply returns an arbitrary option.

PS: show_main_menu function could be renamed in something like launch_action

#!/usr/bin/env python2
# coding=utf-8

def _get_raw_input(prompt):
   while True:
       try:
           return int(raw_input(prompt))
       except ValueError:
           print 'please enter valid input \n'
           continue
       else:
           break

class Foo:
    def menu_post_transaction(self): print 'menu_post_transaction'
    def add_account(self): print "### 1. add_account"
    def login(self): print "### 2. login"

    def show_main_menu(self, get_option=_get_raw_input):
            """shows main menu to user"""
            print '1:create account \t 2:Log in \t any other to exit \n'
            option = get_option('Enter your option: \t')
            if option == 1:
                self.add_account()
                print '\n ***successfully created account***\n'
                self.menu_post_transaction()
            elif option == 2:
                self.login()
            else:
                print 'Thank you for using our services'

if __name__ == "__main__":

    Foo().show_main_menu(lambda _: 1)
    Foo().show_main_menu(lambda _: 2)
    Foo().show_main_menu(lambda _: 3)

This a general good practice independent to the language.

Answer 3

It could something like this:

from unittest.mock import patch
from unittest import TestCase

@patch('yourmodule.show_main_menu', return_value='1')
def test_answer_one(self, input):
    self.assertEqual(show_main_menu(), '\n ***successfully created account***\n')

@patch('yourmodule.show_main_menu', return_value='5')
def test_answer_five(self, input):
    self.assertEqual(show_main_menu(), 'Thank you for using our services')