Refactor to Class

Still test dosen't work
Some code was comment to check if tests will run