Merged Ally's code so that the test pass