Code clean up
slight changes to make tests pass