Fixing logic and test cases accordingly [AMECHE-327]