Mock, import, and coverage

Question

I am writing a small console wrapper so that we can some pretty colours on the console. A cut down version of the module I use is here:

from sys import stderr
try:
    from fabulous.color import fg256
    clog = lambda x: stderr.write(
        str(fg256(63, u'\u26A0 ' + str(x).encode('utf-8') + '\n').as_utf8))
except ImportError:
     _green = '\033[1;32m'
     _default = '\033[1;m'
     clog = lambda x: stderr.write(_green +
                                   u'\u26A0 ' +
                                   str(x).encode('utf-8') +
                                   _default +
                                   '\n')

I can easily write some unit tests using from mycolour improt clog and the right version of clog will be picked up depending on the presence of fabulous or not. However, once the clog lambda is loaded, you cannot unload it. Thus, any attempt to test (via py.test) both paths of the code fails on whichever is the second one.

How can I achieve 100% test coverage for this?

Here's the test script I have:

try:
    import builtins
except ImportError:
    import __builtin__ as builtins
import unittest
from mock import MagicMock
from mock import patch
from mock import call


class ColourFabulousTest(unittest.TestCase):

    def setUp(self):
        pass

    @patch('mycolours.stderr')
    def test_fabulous_clog(self, err):
        from mycolours import clog
        err.write = MagicMock()
        clog('ook')
        self.assertTrue(err.write.called)
        self.assertEqual([call('\x1b[38;5;63m\xe2\x9a\xa0 ook\n\x1b[39m')],
                         err.write.call_args_list)

    def tearDown(self):
        pass


class ColourANSITest(unittest.TestCase):

    def setUp(self):
        self.realimport = builtins.__import__
        def myimport(name, globals, locals, fromlist, level=0):
            if name == 'fabulous.color':
                raise ImportError
            return self.realimport(name, globals, locals, fromlist, level)
        builtins.__import__ = myimport

    def test_ansi_clog(self):
        from mycolours import clog
        builtins.__import__ = self.realimport
        with patch('mycolours.stderr') as err:
            clog('ook')
            self.assertTrue(err.write.called)
            self.assertEqual([call(u'\x1b[1;32m\u26a0 ook\x1b[1;m\n')],
                             err.write.call_args_list)

    def tearDown(self):
        builtins.__import__ = self.realimport

Answer 1

Add this to your tearDown() methods:

try:
    del sys.modules['mycolours']
except KeyError:
    pass

Take a look at this and this more general questions.