Added unit tests to cover additional branches