Major refactor and addition of unit tests