Fix current unit tests
Merge pull request #17 from thescouser89/fix_tests

Fix current unit tests