More unit tests
Minor test refactoring