lrossett / centos / centpkg

Forked from centos/centpkg 2 years ago
Clone
ddb17b
import warnings
ddb17b
ddb17b
__unittest=True
ddb17b
ddb17b
class _AssertWarnsContext(object):
ddb17b
    def __init__(self, warningtype, testcase, msg=''):
ddb17b
        self.warningtype = warningtype
ddb17b
        warnings.filterwarnings('error')
ddb17b
        self.failureException = testcase.failureException
ddb17b
ddb17b
    def __enter__(self):
ddb17b
        return self
ddb17b
ddb17b
    def __exit__(self, exc_type, exc_value, tb):
ddb17b
        if exc_type is None:
ddb17b
            try:
ddb17b
                exc_name = self.warningtype.__name__
ddb17b
            except AttributeError:
ddb17b
                exc_name = str(self.warningtype)
ddb17b
            raise self.failureException(
ddb17b
                "{0} not raised".format(exc_name))
ddb17b
ddb17b
        if not issubclass(exc_type, self.warningtype):
ddb17b
            raise self.failureException('"%s" does not match "%s"' %
ddb17b
                    (self.warningtype.__name__, str(exc_type.__name__)))
ddb17b
ddb17b
        return True
ddb17b
ddb17b
class CatchWarningsMixin(object):
ddb17b
    def assertWarns(self, wrnClass, callableObj=None, *args, **kwargs):
ddb17b
        context = _AssertWarnsContext(wrnClass, self)
ddb17b
        if callableObj is None:
ddb17b
            return context
ddb17b
        with context:
ddb17b
            callableObj(*args, **kwargs)