Fixed failing test after UI was changed