Chore/fix tests

@donnapep @stulees please review. Will need to rebase before merge. Thanks.