Refactor tests

Refactor and add unit tests. It's OK? I think as to small classes are better.