Small fix to test after changing API