Updated tests to achieve full test coverage