Update test.py

small fix