Cleaning up python script. Added some tests.