Fixed code to make tests pass.