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

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

Solution

You can use the catch_warnings context manager. Essentially this allows you to mock the warnings handler, so that you can verify details of the warning. See the official docs for a fuller explanation and sample test code.

import warnings

def fxn():
    warnings.warn("deprecated", DeprecationWarning)

with warnings.catch_warnings(record=True) as w:
    # Cause all warnings to always be triggered.
    warnings.simplefilter("always")
    # Trigger a warning.
    fxn()
    # Verify some things
    assert len(w) == 1
    assert issubclass(w[-1].category, DeprecationWarning)
    assert "deprecated" in str(w[-1].message)