Adding missing unit test coverage