Improve mapping method
Add unit test