Added tests and restructured the code.