Fixed those failing unit tests