Cleanup and refactor unit tests