Add unit tests
