Add unit tests that cover the main functionality