Made some changes to make my tests pass.