Revert "Revert "Fix failing tests""

This reverts commit c2a047cf99f3d020bab12701e8dafed7424673bf.