Some unnittest + Mock examples in Python. Includes examples on how to mock a entire class (ZipFile), mock an Iterator object and how to use side_effect properly

unnittestexample.py

import unittest
import os
 
from zipfile import ZipFile
from mock import MagicMock, patch, Mock, mock_open

# The functions that are tested:
def function_to_test_zipfile(example_arg):
    with ZipFile(example_arg, 'r') as zip_in:
        for input_file in zip_in.infolist():
            with zip_in.open(input_file):
                example_function(example_arg, input_file.filename)


def example_function(example_arg, filename):
        pass

def example_csv_reader(csv_reader):
    counter = 0
    for _ in csv_reader:
        counter += 1
    return counter

def rename_index_file(file_path, new_file_name):
    folder, _ = os.path.split(file_path)
    new_file_name_full = os.path.join(folder, new_file_name)
    n_file = 1
    while os.path.isfile(new_file_name_full):
        new_file_name_full = os.path.join(folder, '%s (%d).csv' %
                                          (new_file_name, n_file))
        n_file += 1
        os.rename(file_path, new_file_name_full)

# The unittests:
class ExampleUnitTest(unittest.TestCase):
    class MockZipf:
        def __init__(self):
                       self.files = [Mock(filename='foo.csv'), Mock(filename='bar.csv')]

        def __iter__(self):
                       return iter(self.files)

        def __enter__(self):
                       return self

        def __exit__(self, exc_type, exc_val, exc_tb):
                       return True

        def infolist(self):
                      return self.files

    # Example unittest that shows how to Mock a class such as ZipFile, and how to assert that a function is called
    # correctly
    @patch('unnittestexample.ZipFile')
    @patch('unnittestexample.example_function')
    def test_function_to_test_calls_example_function(self, example_function_fn, mock_zipfile):
        """Integration test for function_to_test, should call example_function twice with arguments from
        inside the zipfile"""
        mock_open_fn = mock_open(read_data='')
        mock_zipfile.return_value = ExampleUnitTest.MockZipf()
        mock_zipfile.return_value.open = mock_open_fn
 
        with patch('__builtin__.open', mock_open_fn):
            function_to_test_zipfile('example_arg')
 
        example_function_fn.assert_any_call('example_arg', 'foo.csv')
        example_function_fn.assert_any_call('example_arg', 'bar.csv')
        self.assertEqual(2, example_function.call_count)

    # Example unnittest that shows how to mock an Iterator object, such as CSV reader
    def test_example_csv_reader(self):
        """Should count the number of lines in a CSV file"""
        mock_csv_reader = MagicMock()
        mock_csv_reader.__iter__.return_value = iter([['id'],
                                                      ['1'],
                                                      ['1']])
        with patch('__builtin__.open'):
            n_rows = example_csv_reader(mock_csv_reader)
        self.assertEqual(3, n_rows)

    # Example on how to work with side effects. Here os.path.isfile will return True twice, and False once.
    @patch('unnittestexample.os.rename')
    @patch('unnittestexample.os.path.isfile')
    def test_rename_index_file_file_exists(self, os_fn_exists, os_fn_rename):
        """Should call os.rename with the new filename + an added number if file already exists"""
        os_fn_exists.side_effect = [True, True, False]
        rename_index_file('test/foo/bar/test', 'baz')
        os_fn_rename.assert_called_with('test/foo/bar/test',
                                        'test/foo/bar/baz (2).csv')

Good stuff, I like, No more struggle testing ...