Refactor unit tests
quick fixes