Fixing tests break