Add test and fix error