Cleaning up test
Bug fixes after running tests