[Andrew / Augustine] Fixed failing tests