Minor update to tests