Вызвать предупреждение в Python без прерывания программы

Я имею дело с проблемой, как подать предупреждение в Python, не допуская сбоя / остановки / прерывания программы.

Я использую следующую простую функцию, которая только проверяет, передал ли ей пользователь ненулевое число. Если пользователь передает ноль, программа должна предупредить пользователя, но продолжить работу в обычном режиме. Он должен работать как следующий код, но должен использовать класс Warning (), Error () или Exception () вместо того, чтобы выводить предупреждение вручную.

def is_zero(i):
   if i != 0:
     print "OK"
   else:
     print "WARNING: the input is 0!"
   return i

Если я использую приведенный ниже код и передаю 0 функции, программа вылетает и значение никогда не возвращается. Вместо этого я хочу, чтобы программа продолжала работать нормально, и просто сообщала пользователю, что он передал функции 0.

def is_zero(i):
   if i != 0:
     print "OK"
   else:
     raise Warning("the input is 0!")
   return i

Я хочу иметь возможность проверить, было ли выдано предупреждение, тестируя его с помощью unittest. Если я просто распечатаю сообщение, я не смогу проверить его с помощью assertRaises в unittest.