Fix feature problem
Add unit tests