modified code to make tests pass.