Fixed broken tests temporarily