Fixing test cases
Merging with changes from previous commit