Add Test and fix some of the code.