Refactor test suite
Minor changes