Refactor some code. Added new test file