Fix tests for the new API

I have yet to add more tests for that