Added some tests. Performed some refactoring.