Add gracefull stop
Add unit tests & improve some code