Unit test refactor pt 1