Revert "Revert "Fix tests""

This reverts commit 26947ae2d02553c409420a73e38ff884ed1c6634.