[Redisenterprise] az redisenterprise update: Unset capacity and zones while updating sku from enterprise to azure mananaged redis skus as they don't support the properties (#9529)

* Fix Redis Enterprise SKU update and zone configuration issues

* Add examples for redisenterprise update command to fix linting error

* Add flashoptimized in sku check

* remove spaces

* remove spaces

* update helper in aaz registration

* update version

Author: samsaha-ms

src/redisenterprise/HISTORY.rst

@@ -2,6 +2,9 @@
 
 Release History
 ===============
+1.3.1
+- Fixed an issue where updating sku from Azure Cache for Enterprise to Azure Managed Redis SKU was not working as expected.
+
 1.3.0
 - Added a new required property: PublicNetworkAccess for Cluster.
 - Updated the default value of AccessKeysAuthentication property for Database to 'Disabled'.

src/redisenterprise/azext_redisenterprise/aaz/latest/redisenterprise/_update.py

@@ -16,6 +16,9 @@
 )
 class Update(AAZCommand):
     """Update an existing (overwrite/recreate, with potential downtime) cache cluster
+
+    :example: Updates cluster SKU
+        az redisenterprise update --cluster-name "cache1" --sku "ComputeOptimized_X5" --resource-group "rg1"
     """
 
     _aaz_info = {

src/redisenterprise/azext_redisenterprise/commands.py

@@ -16,6 +16,8 @@ def load_command_table(self, _):  # pylint: disable=unused-argument
         g.custom_command('create', 'redisenterprise_create', supports_no_wait=True)
         g.custom_command('list', 'redisenterprise_list')
         g.custom_show_command('show', 'redisenterprise_show')
+        from .custom import RedisEnterpriseUpdate
+        self.command_table["redisenterprise update"] = RedisEnterpriseUpdate(loader=self)
     with self.command_group("redisenterprise database"):
         from .custom import DatabaseFlush, DatabaseCreate, DatabaseDelete, DatabaseExport, DatabaseForceUnlink
         from .custom import DatabaseImport, DatabaseListKey, DatabaseRegenerateKey

src/redisenterprise/azext_redisenterprise/custom.py

@@ -24,13 +24,51 @@
 from .aaz.latest.redisenterprise import List as _ClusterList
 from .aaz.latest.redisenterprise import Show as _ClusterShow
 from .aaz.latest.redisenterprise import Wait as _DatabaseWait
+from .aaz.latest.redisenterprise import Update as _Update
 from azure.cli.core.azclierror import (
     MutuallyExclusiveArgumentError,
 )
 
 logger = get_logger(__name__)
 
 
+class RedisEnterpriseUpdate(_Update):
+
+    def pre_instance_update(self, instance):
+        """Called before the instance is updated"""
+        try:
+            current_sku = str(instance.sku.name) if hasattr(instance, 'sku') and instance.sku else None
+        except (AttributeError, TypeError):
+            current_sku = None
+
+        new_sku = str(self.ctx.args.sku) if self.ctx.args.sku is not None else None
+
+        if new_sku and current_sku and new_sku != current_sku:
+            self._handle_sku_change(current_sku, new_sku, instance)
+
+    def _handle_sku_change(self, current_sku, new_sku, instance):
+        """Handle SKU change logic for capacity and zones"""
+        # Check if changing from Azure Cache for Redis Enterprise to Azure Managed Redis SKU types
+        # that don't support capacity/zones
+        if (current_sku.startswith('Enterprise') and
+            (new_sku.startswith('Balanced_') or
+             new_sku.startswith('ComputeOptimized_') or
+             new_sku.startswith('MemoryOptimized_') or
+             new_sku.startswith('FlashOptimized_'))):
+            # Unset capacity and zones in the instance
+            try:
+                if hasattr(instance, 'sku') and instance.sku and hasattr(instance.sku, 'capacity'):
+                    instance.sku.capacity = None
+            except (AttributeError, TypeError):
+                pass
+
+            try:
+                if hasattr(instance, 'zones'):
+                    instance.zones = None
+            except (AttributeError, TypeError):
+                pass
+
+
 class DatabaseFlush(_DatabaseFlush):
 
     @classmethod

src/redisenterprise/azext_redisenterprise/tests/latest/example_steps.py

@@ -97,6 +97,21 @@ def step_create(test, checks=None, cache_num=1):
                  '--port 10000 '
                  '--resource-group "{rg}"',
                  checks=checks)
+    elif test.kwargs.get('sku-update'):
+        test.cmd('az redisenterprise create '
+                 '--cluster-name "{cluster}" '
+                 '--sku "{initial_sku}" '
+                 '--location "centraluseuap" '
+                 '--tags tag1="value1" '
+                 '--minimum-tls-version "1.2" '
+                 '--client-protocol "Encrypted" '
+                 '--clustering-policy "EnterpriseCluster" '
+                 '--public-network-access "Enabled" '
+                 '--access-keys-auth Enabled '
+                 '--eviction-policy "NoEviction" '
+                 '--port 10000 '
+                 '--resource-group "{rg}"',
+                 checks=checks)
     else:
         test.cmd('az redisenterprise create '
                  '--cluster-name "{cluster}" '
@@ -165,6 +180,16 @@ def step_delete(test, checks=None):
                  '--resource-group "{rg}"',
                  checks=checks)
 
+def step_update(test, checks=None):
+    if checks is None:
+        checks = []
+    if test.kwargs.get('sku-update'):
+        test.cmd('az redisenterprise update '
+                 '--cluster-name "{cluster}" '
+                 '--sku "{new_sku}" '
+                 '--resource-group "{rg}"',
+                 checks=checks)
+
 def step_database_update(test, checks=None):
     if checks is None:
         checks = []

src/redisenterprise/azext_redisenterprise/tests/latest/recordings/test_redisenterprise_scenario7.yaml

@@ -20,6 +20,7 @@
 from .example_steps import step_database_delete
 from .example_steps import step_database_force_unlink
 from .example_steps import step_database_update
+from .example_steps import step_update
 from .example_steps import step_database_access_policy_assignment_create
 from .example_steps import step_database_access_policy_assignment_list
 from .example_steps import step_database_access_policy_assignment_delete
@@ -588,3 +589,96 @@ def test_redisenterprise_scenario6(self, rg):
         call_scenario6(self, rg)
         calc_coverage(__file__)
         raise_if()
+
+
+# Env setup_scenario7
+@try_manual
+def setup_scenario7(test):
+    pass
+
+
+# Env cleanup_scenario7
+@try_manual
+def cleanup_scenario7(test):
+    pass
+
+
+# Testcase: scenario7 - SKU Update Testing for RedisEnterpriseUpdate
+def call_scenario7(test, rg):
+    """Test scenario specifically for SKU update functionality with RedisEnterpriseUpdate class"""
+    setup_scenario7(test)
+    
+    # Create initial cluster
+    step_create(test, checks=[
+        test.check("name", "default"),
+        test.check("resourceGroup", "{rg}"),
+        test.check("clientProtocol", "Encrypted"),
+        test.check("clusteringPolicy", "EnterpriseCluster"),
+        test.check("evictionPolicy", "NoEviction"),
+        test.check("port", 10000),
+        test.check("provisioningState", "Succeeded"),
+        test.check("resourceState", "Running"),
+        test.check("type", "Microsoft.Cache/redisEnterprise/databases")
+    ])
+    
+    # Verify initial cluster with Enterprise SKU has capacity and zones
+    step_show(test, checks=[
+        test.check("name", "{cluster}"),
+        test.check("resourceGroup", "{rg}"),
+        test.check("location", "Central US EUAP"),
+        test.check("sku.name", "Balanced_B5"),
+        test.check("sku.capacity", None),
+        test.check("zones", None),
+        test.check("provisioningState", "Succeeded"),
+        test.check("resourceState", "Running"),
+        test.check("type", "Microsoft.Cache/redisEnterprise"),
+        test.check("databases[0].name", "default")
+    ])
+    
+    step_update(test, checks=[
+        test.check("name", "{cluster}"),
+        test.check("resourceGroup", "{rg}"),
+        test.check("sku.name", "ComputeOptimized_X5"),
+        test.check("sku.capacity", None),
+        test.check("zones", None),
+        test.check("provisioningState", "Succeeded"),
+        test.check("resourceState", "Running"),
+        test.check("type", "Microsoft.Cache/redisEnterprise")
+    ])
+    
+    # Verify the updated cluster state
+    step_show(test, checks=[
+        test.check("name", "{cluster}"),
+        test.check("resourceGroup", "{rg}"),
+        test.check("sku.name", "ComputeOptimized_X5"),
+        test.check("sku.capacity", None),
+        test.check("zones", None),
+        test.check("provisioningState", "Succeeded"),
+        test.check("resourceState", "Running"),
+        test.check("type", "Microsoft.Cache/redisEnterprise")
+    ])
+    
+    step_delete(test, checks=[])
+    cleanup_scenario7(test)
+
+
+# Test class for scenario7 - SKU Update Testing
+class Redisenterprisescenario7Test(ScenarioTest):
+    
+    def __init__(self, *args, **kwargs):
+        super(Redisenterprisescenario7Test, self).__init__(*args, **kwargs)
+
+        self.kwargs.update({
+            'cluster': self.create_random_name(prefix='clitest-cache7-', length=21),
+            'sku-update': True,
+            'initial_sku': 'Balanced_B5',
+            'new_sku': 'ComputeOptimized_X5'
+        })
+
+    @AllowLargeResponse(size_kb=9999)
+    @ResourceGroupPreparer(name_prefix='clitest-redisenterprise-rg7-', key='rg', parameter_name='rg',
+                           location='centraluseuap', random_name_length=34)
+    def test_redisenterprise_scenario7(self, rg):
+        call_scenario7(self, rg)
+        calc_coverage(__file__)
+        raise_if()

src/redisenterprise/azext_redisenterprise/tests/latest/test_demo.py

@@ -10,7 +10,7 @@
 
 
 # HISTORY.rst entry.
-VERSION = '1.3.0'
+VERSION = '1.3.1'
 
 # The full list of classifiers is available at
 # https://pypi.python.org/pypi?%3Aaction=list_classifiers