How to test with Python's unittest that a warning has been thrown?

Viewed 27117

I have a following function in Python and I want to test with unittest that if the function gets 0 as argument, it throws a warning. I already tried assertRaises, but since I don't raise the warning, that doesn't work.

def isZero(i):
    if i != 0:
        print "OK"
    else:
        warning = Warning("the input is 0!") 
        print warning
    return i
6 Answers

One problem with the warnings.catch_warnings approach is that warnings produced in different tests can interact in strange ways through global state kept in __warningregistry__ attributes.

To address this, we should clear the __warningregistry__ attribute of every module before every test that checks warnings.

class MyTest(unittest.TestCase):

  def setUp(self):
    # The __warningregistry__'s need to be in a pristine state for tests
    # to work properly.
    for v in sys.modules.values():
      if getattr(v, '__warningregistry__', None):
        v.__warningregistry__ = {}

  def test_something(self):
    with warnings.catch_warnings(record=True) as w:
      warnings.simplefilter("always", MySpecialWarning)
      ...
      self.assertEqual(len(w), 1)
      self.assertIsInstance(w[0].message, MySpecialWarning)

This is how Python 3's assertWarns() method is implemented.

Building off the answer from @ire_and_curses,

class AssertWarns(warnings.catch_warnings):
    """A Python 2 compatible version of `unittest.TestCase.assertWarns`."""
    def __init__(self, test_case, warning_type):
        self.test_case = test_case
        self.warning_type = warning_type
        self.log = None
        super(AssertWarns, self).__init__(record=True, module=None)

    def __enter__(self):
        self.log = super(AssertWarns, self).__enter__()
        return self.log

    def __exit__(self, *exc_info):
        super(AssertWarns, self).__exit__(*exc_info)
        self.test_case.assertEqual(type(self.log[0]), self.warning_type)

This can be called similarly to unittest.TestCase.assertWarns:

with AssertWarns(self, warnings.WarningMessage):
    warnings.warn('test warning!') 

where self is a unittest.TestCase.

Related