Fixed test to use new stuff in API