Tests adjusted for new API