Final changes to make all tests pass :)