got tests working after refactor