Added unit tests

Addresses issue #4