Finished first test and refactoring