Added main class
Fixed failing tests