Get test coverage back up to 100%