Add test and make it pass