More clean up after big refactor