Almost done fixing tests