Adding tests and fixing edge case.