From f1ea4cad98f5abc17f6afa82f6d38a947a4c9d98 Mon Sep 17 00:00:00 2001
From: deveshdama <87668846+deveshdama@users.noreply.github.com>
Date: Wed, 6 Mar 2024 23:54:17 -0800
Subject: [PATCH] {AKS} `az aks mesh get-upgrades` fails with a traceback if
 ASM addon is disabled (#7353)

---
 src/aks-preview/HISTORY.rst                 |  3 ++-
 src/aks-preview/azext_aks_preview/custom.py | 28 ++++++++++++++-------
 src/aks-preview/setup.py                    |  2 +-
 3 files changed, 22 insertions(+), 11 deletions(-)

diff --git a/src/aks-preview/HISTORY.rst b/src/aks-preview/HISTORY.rst
index e830d558c83..4a2aab5519d 100644
--- a/src/aks-preview/HISTORY.rst
+++ b/src/aks-preview/HISTORY.rst
@@ -9,9 +9,10 @@ If there is no rush to release a new version, please just add a description of t
 
 To release a new version, please select a new version number (usually plus 1 to last patch version, X.Y.Z -> Major.Minor.Patch, more details in `\doc <https://semver.org/>`_), and then add a new section named as the new version number in this file, the content should include the new modifications and everything from the *Pending* section. Finally, update the `VERSION` variable in `setup.py` with this new version number.
 
-Pending
+2.0.0b3
 +++++++
 * Add parameter to set revision `--revision` for the Azure Service Mesh addon while creating AKS cluster.
+* Fix for `az aks mesh get-upgrades` command panic response when ASM addon is not enabled.
 
 2.0.0b2
 +++++++
diff --git a/src/aks-preview/azext_aks_preview/custom.py b/src/aks-preview/azext_aks_preview/custom.py
index a8ca1a50602..81059bf24cd 100644
--- a/src/aks-preview/azext_aks_preview/custom.py
+++ b/src/aks-preview/azext_aks_preview/custom.py
@@ -2947,6 +2947,17 @@ def aks_mesh_get_revisions(
     return None
 
 
+def check_iterator(iterator):
+    import itertools
+    try:
+        first = next(iterator)
+    except StopIteration:   # iterator is empty
+        return True, iterator
+    except TypeError:       # iterator is not iterable, e.g. None
+        return True, iterator
+    return False, itertools.chain([first], iterator)
+
+
 def aks_mesh_get_upgrades(
         cmd,
         client,
@@ -2954,15 +2965,14 @@ def aks_mesh_get_upgrades(
         name
 ):
     upgradeProfiles = client.list_mesh_upgrade_profiles(resource_group_name, name)
-    # 'upgradeProfiles' is an ItemPaged object
-    upgrades = []
-    # Iterate over items within pages
-    for page in upgradeProfiles.by_page():
-        for item in page:
-            upgrades.append(item)
-
-    if upgrades:
-        return upgrades[0].properties
+    is_empty, upgradeProfiles = check_iterator(upgradeProfiles)
+    if is_empty:
+        logger.warning("No mesh upgrade profiles found for the cluster '%s' " +
+                       "in the resource group '%s'.", name, resource_group_name)
+        return None
+    upgrade = next(upgradeProfiles, None)
+    if upgrade:
+        return upgrade.properties
     return None
 
 
diff --git a/src/aks-preview/setup.py b/src/aks-preview/setup.py
index e9285ae2689..58f415b7cbe 100644
--- a/src/aks-preview/setup.py
+++ b/src/aks-preview/setup.py
@@ -9,7 +9,7 @@
 
 from setuptools import setup, find_packages
 
-VERSION = "2.0.0b2"
+VERSION = "2.0.0b3"
 
 CLASSIFIERS = [
     "Development Status :: 4 - Beta",