Finally fixing tests