[CLEANUP] Streamline code in UpgradeWizard