[Upgrade] Fix 16952: Fix ImportError after upgrade (#17314)

src/azure-cli/azure/cli/command_modules/util/custom.py

@@ -141,21 +141,24 @@ def upgrade_version(cmd, update_all=None, yes=None):  # pylint: disable=too-many
         telemetry.set_failure(err_msg)
         sys.exit(exit_code)
 
-    import azure.cli.core
+    # Avoid using python modules directly as they may have been changed due to upgrade.
+    # If you do need to use them, you may need to reload them and their dependent modules.
+    # Otherwise you may have such issue https://github.com/Azure/azure-cli/issues/16952
     import importlib
-    importlib.reload(azure.cli.core)
-    new_version = azure.cli.core.__version__
+    import json
+    importlib.reload(subprocess)
+    importlib.reload(json)
+
+    version_result = subprocess.check_output(['az', 'version', '-o', 'json'], shell=platform.system() == 'Windows')
+    version_json = json.loads(version_result)
+    new_version = version_json['azure-cli-core']
 
     if update_cli and new_version == local_version:
         err_msg = "CLI upgrade failed or aborted."
         logger.warning(err_msg)
         telemetry.set_failure(err_msg)
         sys.exit(1)
 
-    # Python is reinstalled in another versioned directory with Homebrew, subprocess needs to be reloaded
-    if installer == 'HOMEBREW' and exts:
-        importlib.reload(subprocess)
-
     if exts:
         logger.warning("Upgrading extensions")
     for ext_name in exts: