Added more unit tests and made some refactors.