Fixed failing tests from changes