Cleanup stuff

@raghavgautam  were you able to verify the above change ? last time changes were made to this method it broke lot of test, so i wanted to be sure.