Along the same lines - use doctests maybe?

    if __name__ == "__main__":
        import doctest
        doctest.testmod()

https://docs.python.org/3.5/library/doctest.html