Refacted classes and updated tests