Small updates to tests