Added unit tests, fixed a bug.