Updated tests for changed cherry-picked to this branch