updating unit tests to work post refactor