Search code examples
pythonunit-testingmockingmonkeypatching

Testing branches in module-level code in python


We have a python module that a caller can use to run some utility commands on Mac OS X. The path to the commands and their usage differ between versions of the OS and our module is intended to hide that from the user. We determine the version of the OS once when the module is imported, like so (helper.py, simplified to demonstrate the point):

import platform
from distutils.version import StrictVersion

print('Importing helper module...')

os_version, _, _ = platform.mac_ver()

if StrictVersion(os_version) < StrictVersion('10.10'):
    utility1 = 'utility 1 for 10.9 and below'
else:
    utility1 = 'utility 1 for 10.10 and above'

def run_utility1():
    return 'Running ' + utility1

def run_utility2():
    # ...
    pass

# ... more cool functions ...

Now we'd like to add tests to this module. Specifically, we'd like to make sure that the correct utility runs for all versions of OS X. The way I thought of was to patch platform.mac_ver() in different tests to return a different OS version, and assert that we're running the right utility. Like so:

import mock
import unittest

class HelperTests(unittest.TestCase):
    def test_10_9_utility_is_correct(self):
        with mock.patch('platform.mac_ver', return_value=('10.9', 'foo', 'foo')):
            import helper
            result = helper.run_utility1()
            print(result)
            assert result == 'Running utility 1 for 10.9 and below'

    def test_10_10_utility_is_correct(self):
        with mock.patch('platform.mac_ver', return_value=('10.10', 'foo', 'foo')):
            import helper
            result = helper.run_utility1()
            print(result)
            assert result == 'Running utility 1 for 10.10 and above'

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

But this results in:

Testing started at 12:16 PM ...
Importing helper module...
Running utility 1 for 10.10 and above
Running utility 1 for 10.10 and above

Process finished with exit code 0

Failure
Traceback (most recent call last):
  File "helper_test.py", line 13, in test_10_9_utility_is_correct
    assert result == 'Running utility 1 for 10.9 and below'
AssertionError

It seems that test_10_10_utility_is_correct is running first (is this due to alphabetical ordering of methods in the test harness?), patching mac_ver() to return 10.10 and then importing helper. When test_10_9_utility_is_correct runs helper isn't imported again, so it fails since it thinks it's on 10.10.

I understand that python doesn't import a module twice, and that's awesome. But does that mean we can't exercise branches in module-level code in tests, since it will only run once? If there is a way to do so, how?

I've considered wrapping the module-level OS version checking code in a function. That would make it easy to mock but then all other functions would have to call it first, which seems unnecessary since the OS version is not likely to change between calls. I've also considered moving each test method into its own test module, which would cause helper to get imported multiple times, but that seems pretty clunky. Is there another way to exercise the two branches given in helper.py?


Solution

  • How about putting your initialisation code in a function ?

    You will have to declare your global variable like this :

    import platform
    from distutils.version import StrictVersion
    
    utility1 = None
    def init_module():
        global utility1 # declare it global to modify it
        print('Importing helper module...')
    
        os_version, _, _ = platform.mac_ver()
    
        if StrictVersion(os_version) < StrictVersion('10.10'):
            utility1 = 'utility 1 for 10.9 and below'
        else:
            utility1 = 'utility 1 for 10.10 and above'
    
    init_module() #Call the init function when importing the module
    
    def run_utility1():
        return 'Running ' + utility1
    

    Then at each new test, you can call the init function :

    import mock
    import unittest
    import helper
    
    class HelperTests(unittest.TestCase):
        def test_10_9_utility_is_correct(self):
            with mock.patch('platform.mac_ver', return_value=('10.9', 'foo', 'foo')):
                helper.init_module()
                result = helper.run_utility1()
                print(result)
                assert result == 'Running utility 1 for 10.9 and below'
    
        def test_10_10_utility_is_correct(self):
            with mock.patch('platform.mac_ver', return_value=('10.10', 'foo', 'foo')):
                helper.init_module()
                result = helper.run_utility1()
                print(result)
                assert result == 'Running utility 1 for 10.10 and above'
    
    if __name__ == '__main__':
        unittest.main()