Fixing test after API change