Fixing remaining tests