Adding tests and coverage