Fix remaining failing tests (hopefully)