Fix test failures caused by previous commit