Adapt test to new design