Added unit tests (finally)