Fix failing tests (temporary fix for 6.0.1)