Make tests compatible with both py3 and py2

Signed-off-by: Vasudev Kamath <kamathvasudev@gmail.com>