Fix test failing on CI