Add new tests, refacatored code
Merge