Updated logic to fix test

Added few tests