Added if statements to handle wrong test cases