ADD - Unit tests, fixed some bugs due to failing unit tests.