from __future__ import print_function

class bcolors: # thanks to https://stackoverflow.com/questions/287871/how-do-i-print-colored-text-to-the-terminal

HEADER = '\033[95m'

OKBLUE = '\033[94m'

OKCYAN = '\033[96m'

OKGREEN = '\033[92m'

WARNING = '\033[93m'

FAIL = '\033[91m'

ENDC = '\033[0m'

BOLD = '\033[1m'

UNDERLINE = '\033[4m'

class AssertException(Exception):

pass

def format_message(message):

# return message.replace("\n", "<:LF:>")

return message

def display(type, message, label="", mode=""):

if type == "PASSED":

color = bcolors.OKGREEN

elif type == "FAILED":

color = bcolors.WARNING

elif type == "ERROR":

color = bcolors.FAIL

elif type == "COMPLETEDIN":

color = bcolors.OKCYAN

else:

color = bcolors.UNDERLINE

print("\n{4}<{0}:{1}:{2}>{3}{5}".format(

type.upper(), mode.upper(), label, format_message(message), color, bcolors.ENDC)

, flush=True)

def expect(passed=None, message=None, allow_raise=False):

if passed:

display('PASSED', 'Test Passed')

else:

message = message or "Value is not what was expected"

display('FAILED', message)

if allow_raise:

raise AssertException(message)

def assert_equals(actual, expected, message=None, allow_raise=False):

equals_msg = "{0} should equal {1}".format(repr(actual), repr(expected))

if message is None:

message = equals_msg

else:

message += ": " + equals_msg

expect(actual == expected, message, allow_raise)

def assert_not_equals(actual, expected, message=None, allow_raise=False):

r_actual, r_expected = repr(actual), repr(expected)

equals_msg = "{0} should not equal {1}".format(r_actual, r_expected)

if message is None:

message = equals_msg

else:

message += ": " + equals_msg

expect(not (actual == expected), message, allow_raise)

def expect_error(message, function, exception=Exception):

passed = False

try:

function()

except exception:

passed = True

except Exception as e:

message = "{}: {} should be {}".format(message or "Unexpected exception", repr(e), repr(exception))

expect(passed, message)

def expect_no_error(message, function, exception=BaseException):

try:

function()

except exception as e:

fail("{}: {}".format(message or "Unexpected exception", repr(e)))

return

except:

pass

pass_()

def pass_(): expect(True)

def fail(message): expect(False, message)

def assert_approx_equals(

actual, expected, margin=1e-9, message=None, allow_raise=False):

msg = "{0} should be close to {1} with absolute or relative margin of {2}"

equals_msg = msg.format(repr(actual), repr(expected), repr(margin))

if message is None:

message = equals_msg

else:

message += ": " + equals_msg

div = max(abs(actual), abs(expected), 1)

expect(abs((actual - expected) / div) < margin, message, allow_raise)

'''

Usage:

@describe('describe text')

def describe1():

@it('it text')

def it1():

# some test cases...

'''

def _timed_block_factory(opening_text):

from timeit import default_timer as timer

from traceback import format_exception

from sys import exc_info

def _timed_block_decorator(s, before=None, after=None):

display(opening_text, s)

def wrapper(func):

if callable(before):

before()

time = timer()

try:

func()

except AssertionError as e:

display('FAILED', str(e))

except Exception:

fail('Unexpected exception raised')

tb_str = ''.join(format_exception(*exc_info()))

display('ERROR', tb_str)

display('COMPLETEDIN', '{:.2f}'.format((timer() - time) * 1000))

if callable(after):

after()

return wrapper

return _timed_block_decorator

describe = _timed_block_factory('DESCRIBE')

it = _timed_block_factory('IT')

'''

Timeout utility

Usage:

@timeout(sec)

def some_tests():

any code block...

Note: Timeout value can be a float.

'''

def timeout(sec):

def wrapper(func):

from multiprocessing import Process

msg = 'Should not throw any exceptions inside timeout'

def wrapped():

expect_no_error(msg, func)

process = Process(target=wrapped)

process.start()

process.join(sec)

if process.is_alive():

fail('Exceeded time limit of {:.3f} seconds'.format(sec))

process.terminate()

process.join()

return wrapper