Added unit tests. Fixed some issues.