Refactoring code toward new API