Fixes for refactor bugs