unit test fix after refactor