Added unit test, fixed bug.