Revert "Revert "Maybe fix unit tests""

This reverts commit 2ee1973f5506d24dd9891f58576149c3922c3ef0.