Adding main unit tests
Fix code after writing unit tests