Search code examples
pythonunit-testingtestingmockingtraceback

Faking a traceback in Python


I'm writing a test runner. I have an object that can catch and store exceptions, which will be formatted as a string later as part of the test failure report. I'm trying to unit-test the procedure that formats the exception.

In my test setup, I don't want to actually throw an exception for my object to catch, mainly because it means that the traceback won't be predictable. (If the file changes length, the line numbers in the traceback will change.)

How can I attach a fake traceback to an exception, so that I can make assertions about the way it's formatted? Is this even possible? I'm using Python 3.3.

Simplified example:

class ExceptionCatcher(object):
    def __init__(self, function_to_try):
        self.f = function_to_try
        self.exception = None
    def try_run(self):
        try:
            self.f()
        except Exception as e:
            self.exception = e

def format_exception_catcher(catcher):
    pass
    # No implementation yet - I'm doing TDD.
    # This'll probably use the 'traceback' module to stringify catcher.exception


class TestFormattingExceptions(unittest.TestCase):
    def test_formatting(self):
        catcher = ExceptionCatcher(None)
        catcher.exception = ValueError("Oh no")

        # do something to catcher.exception so that it has a traceback?

        output_str = format_exception_catcher(catcher)
        self.assertEquals(output_str,
"""Traceback (most recent call last):
  File "nonexistent_file.py", line 100, in nonexistent_function
    raise ValueError("Oh no")
ValueError: Oh no
""")

Solution

  • Reading the source of traceback.py pointed me in the right direction. Here's my hacky solution, which involves faking the frame and code objects which the traceback would normally hold references to.

    import traceback
    
    class FakeCode(object):
        def __init__(self, co_filename, co_name):
            self.co_filename = co_filename
            self.co_name = co_name
    
    
    class FakeFrame(object):
        def __init__(self, f_code, f_globals):
            self.f_code = f_code
            self.f_globals = f_globals
    
    
    class FakeTraceback(object):
        def __init__(self, frames, line_nums):
            if len(frames) != len(line_nums):
                raise ValueError("Ya messed up!")
            self._frames = frames
            self._line_nums = line_nums
            self.tb_frame = frames[0]
            self.tb_lineno = line_nums[0]
    
        @property
        def tb_next(self):
            if len(self._frames) > 1:
                return FakeTraceback(self._frames[1:], self._line_nums[1:])
    
    
    class FakeException(Exception):
        def __init__(self, *args, **kwargs):
            self._tb = None
            super().__init__(*args, **kwargs)
    
        @property
        def __traceback__(self):
            return self._tb
    
        @__traceback__.setter
        def __traceback__(self, value):
            self._tb = value
    
        def with_traceback(self, value):
            self._tb = value
            return self
    
    
    code1 = FakeCode("made_up_filename.py", "non_existent_function")
    code2 = FakeCode("another_non_existent_file.py", "another_non_existent_method")
    frame1 = FakeFrame(code1, {})
    frame2 = FakeFrame(code2, {})
    tb = FakeTraceback([frame1, frame2], [1,3])
    exc = FakeException("yo").with_traceback(tb)
    
    print(''.join(traceback.format_exception(FakeException, exc, tb)))
    # Traceback (most recent call last):
    #   File "made_up_filename.py", line 1, in non_existent_function
    #   File "another_non_existent_file.py", line 3, in another_non_existent_method
    # FakeException: yo
    

    Thanks to @User for providing FakeException, which is necessary because real exceptions type-check the argument to with_traceback().

    This version does have a few limitations:

    • It doesn't print the lines of code for each stack frame, as a real traceback would, because format_exception goes off to look for the real file that the code came from (which doesn't exist in our case). If you want to make this work, you need to insert fake data into linecache's cache (because traceback uses linecache to get hold of the source code), per @User's answer below.

    • You also can't actually raise exc and expect the fake traceback to survive.

    • More generally, if you have client code that traverses tracebacks in a different manner than traceback does (such as much of the inspect module), these fakes probably won't work. You'd need to add whatever extra attributes the client code expects.

    These limitations are fine for my purposes - I'm just using it as a test double for code that calls traceback - but if you want to do more involved traceback manipulation, it looks like you might have to go down to the C level.