Update test.py

Cleaned it up