adding unit tests and some refactoring
update diagram