Made the code more robust
Merge branch 'master' of https://github.com/nwotero/cs2223