Refactored tests and classes

Now logically separated classes and some methods to make tests more
readable