Added tests and removed unnecessary code.