Fixed a couple problems with test.py