add tests and fix the code

: