Improve/Fix test cases