Implement unit tests

Implement some unit tests to improve test coverage.