Added test cases and updated functions logic