Refactor tests and pass them