Updated tests for full test coverage