finished code refactor / added test