Re-factored code and implemented unit tests. Fixes #1 and closes #10