Fixed some problems to make it works with the new API

Change-Id: I30b37e2e9e4eff8bd4d4be722a1c31ce566b1b5b