Add unit tests
Refactor code
Improve code
Merge pull request #3 from Nuke928/master

Add basic unit tests