Update test.py

Thanks Gabriel for final touch up, now code is working fine in terminal :)