Fix screenshot tests