Adapt test to recent changes in code