Few more fixes to tests broken during merge