Add test case
Improve if else statement.