Updated unit tests according to code review.