Added few bug fixes
Fixed the tests approach