Simple contextmanager to test for flask.abort()

test.py

'''
    A context manager that checks for an exception to be raised
    And an example of it being used to check that flask.abort() is called
    
    Consists of
     * Method for getting a dictionary of {http status code : werkzeug exception}
     * Context manager for asserting something is raised (AssertRaises)
     * unittest2.TestCase subclass that exposes API to access AssertRaises
    
    AssertRaises is similar to what unittest2.assertRaises uses (unittest2.case._AssertRaisesContext)
    Except:
     * Only propogates AssertionError
     * Will complain if unexpected exception is not the same as expected exception
     * Will match exception exactly, instead of a sublcass of it
'''

from werkzeug import exceptions as werkzeug_exceptions
import unittest2 as unittest
import contextlib
import flask
import re

def make_exception_dict():
    '''
        Create dictionary of {code:exception}
        By looking at the exceptions on werkzeug.exceptions
    '''
    exceptions = {}
    for name in dir(werkzeug_exceptions):
        attr = getattr(werkzeug_exceptions, name)
        if isinstance(attr, type) and hasattr(attr, 'code') and attr.code:
            exceptions[str(attr.code)] = attr
    return exceptions

class AssertRaises(object):
    def __init__(self, testcase, expected, regex=None):
        self.regex = regex        
        self.testcase = testcase
        self.expected = expected
    
    def __enter__(self):
        '''Nothing to do on entry'''
        return self
    
    def expected_name(self):
        '''Determine name of expected exception'''
        if hasattr(self.expected, '__name__'):
            return self.expected.__name__
        else:
            return str(self.expected)
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        '''
            Test that we received the expected exception
            * Complain if nothing was raised
            * Propogate AssertionError
            * Complain if regex for message is provided and re.search(regex, exc_value) fails
        '''
        if not exc_type:
            # Didn't get an exception.... :(
            assert False, "'%s' not raised" % self.expected_name()
        
        if exc_type is AssertionError:
            # Make contextmanager reraise the error
            return False
        
        # Assert exception raised is expected
        self.testcase.assertIs(exc_type, self.expected)
        
        # Make sure msg matches regex provided
        if self.regex:
            if not exc_val:
                assert False, "Expected exception to have a message matching %s" % self.regex
            
            failure_message = "Expected exception to contain '%s': Found '%s' instead" % (self.regex, exc_val)
            assert re.search(self.regex, str(exc_val)), failure_message
        
        # Tell contextmanager to silence the error
        return True

class Base(unittest.TestCase):
    exceptions = make_exception_dict()
    
    def assertAborts(self, code):
        expected = self.exceptions[str(code)]
        return AssertRaises(self, expected)
    
    def assertRaises(self, exception, msg):
        return AssertRaises(self, exception, msg)

class Tester(Base):
    def test_things(self):
        '''passes'''
        a = 1 + 1
        with self.assertAborts(401):
            flask.abort(401)

    def test_things2(self):
        """Fails::
        
        FAIL: test:Tester.test_things2
          test.py  # test_things2
            pass
          test.py  # __exit__
            assert False, "'%s' not raised" % self.expected_name()
        AssertionError: 'Unauthorized' not raised
        """
        a = 1 + 1
        with self.assertAborts(401):
            pass

    def test_things3(self):
        """Fails::
        
        FAIL: test:Tester.test_things3
          test.py  # test_things3
            flask.abort(400)
          test.py  # __exit__
            self.testcase.assertIs(exc_type, self.expected)
        AssertionError: <class 'werkzeug.exceptions.ClientDisconnected'> is not <class 'werkzeug.exceptions.Unauthorized'>
        """
        a = 1 + 1
        with self.assertAborts(401):
            flask.abort(400)
    
    def test_things4(self):
        """passes"""
        with self.assertRaises(Exception, "blah and things"):
            raise Exception, "blah and things"
    
    def test_things5(self):
        """Fails::
        FAIL: test:Tester.test_things5
          test.py  # test_things5
            raise Exception, "asdfsadfsdaf sadf"
          test.py  # __exit__
            assert re.search(self.regex, str(exc_val)), failure_message
        AssertionError: Expected exception to contain 'blah and things': Found 'asdfsadfsdaf sadf' instead
        """
        with self.assertRaises(Exception, "blah and things"):
            raise Exception, "asdfsadfsdaf sadf"
    
    def test_things6(self):
        """Fails::
        
        FAIL: test:Tester.test_things6
          test.py  # test_things6
            raise ValueError, "blah and things"
          test.py  # __exit__
            self.testcase.assertIs(exc_type, self.expected)
        AssertionError: <type 'exceptions.ValueError'> is not <type 'exceptions.Exception'>
        """
        with self.assertRaises(Exception, "blah and things"):
            raise ValueError, "blah and things"