changed some code to make tests pass