Unitary test

Adding tests and cleaning some code