Work on test suite - refactoring and fixing bugs