{Compute} az vmss set-orchestration-service-state: Migrate command to aaz-based implementation

src/azure-cli/azure/cli/command_modules/vm/_params.py

@@ -29,7 +29,8 @@
 
 from azure.cli.command_modules.monitor.validators import validate_metric_dimension
 from azure.cli.command_modules.monitor.actions import get_period_type
-from ._vm_utils import RebootSetting, VMGuestPatchClassificationWindows, VMGuestPatchClassificationLinux
+from ._vm_utils import (RebootSetting, VMGuestPatchClassificationWindows, VMGuestPatchClassificationLinux,
+                        OrchestrationServiceNames, OrchestrationServiceStateAction)
 
 
 # pylint: disable=too-many-statements, too-many-branches, too-many-locals, too-many-lines
@@ -40,7 +41,6 @@ def load_arguments(self, _):
     UpgradeMode, CachingTypes, OperatingSystemTypes = self.get_models('UpgradeMode', 'CachingTypes', 'OperatingSystemTypes')
     HyperVGenerationTypes = self.get_models('HyperVGenerationTypes')
     DedicatedHostLicenseTypes = self.get_models('DedicatedHostLicenseTypes')
-    OrchestrationServiceNames, OrchestrationServiceStateAction = self.get_models('OrchestrationServiceNames', 'OrchestrationServiceStateAction', operation_group='virtual_machine_scale_sets')
     ReplicationMode = self.get_models('ReplicationMode', operation_group='gallery_image_versions')
     DiskControllerTypes = self.get_models('DiskControllerTypes', operation_group='virtual_machines')
 

src/azure-cli/azure/cli/command_modules/vm/_vm_utils.py

@@ -817,3 +817,12 @@ class UpgradeMode(Enum):
     AUTOMATIC = 'Automatic'
     MANUAL = 'Manual'
     ROLLING = 'Rolling'
+
+
+class OrchestrationServiceNames(Enum):
+    AUTOMATIC_REPAIRS = 'AutomaticRepairs'
+
+
+class OrchestrationServiceStateAction(Enum):
+    RESUME = 'Resume'
+    SUSPEND = 'Suspend'

src/azure-cli/azure/cli/command_modules/vm/aaz/latest/vmss/__init__.py

@@ -22,6 +22,7 @@
 from ._perform_maintenance import *
 from ._reimage import *
 from ._reimageall import *
+from ._set_orchestration_service_state import *
 from ._show import *
 from ._simulate_eviction import *
 from ._start import *

src/azure-cli/azure/cli/command_modules/vm/aaz/latest/vmss/_set_orchestration_service_state.py

@@ -0,0 +1,183 @@
+# --------------------------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See License.txt in the project root for license information.
+#
+# Code generated by aaz-dev-tools
+# --------------------------------------------------------------------------------------------
+
+# pylint: skip-file
+# flake8: noqa
+
+from azure.cli.core.aaz import *
+
+
+class SetOrchestrationServiceState(AAZCommand):
+    """Changes ServiceState property for a given service
+    """
+
+    _aaz_info = {
+        "version": "2024-11-01",
+        "resources": [
+            ["mgmt-plane", "/subscriptions/{}/resourcegroups/{}/providers/microsoft.compute/virtualmachinescalesets/{}/setorchestrationservicestate", "2024-11-01"],
+        ]
+    }
+
+    AZ_SUPPORT_NO_WAIT = True
+
+    def _handler(self, command_args):
+        super()._handler(command_args)
+        return self.build_lro_poller(self._execute_operations, None)
+
+    _args_schema = None
+
+    @classmethod
+    def _build_arguments_schema(cls, *args, **kwargs):
+        if cls._args_schema is not None:
+            return cls._args_schema
+        cls._args_schema = super()._build_arguments_schema(*args, **kwargs)
+
+        # define Arg Group ""
+
+        _args_schema = cls._args_schema
+        _args_schema.resource_group = AAZResourceGroupNameArg(
+            required=True,
+        )
+        _args_schema.vm_scale_set_name = AAZStrArg(
+            options=["--vm-scale-set-name"],
+            help="The name of the VM scale set.",
+            required=True,
+            id_part="name",
+        )
+
+        # define Arg Group "Parameters"
+
+        _args_schema = cls._args_schema
+        _args_schema.action = AAZStrArg(
+            options=["--action"],
+            arg_group="Parameters",
+            help="The action to be performed.",
+            required=True,
+            enum={"Resume": "Resume", "Suspend": "Suspend"},
+        )
+        _args_schema.service_name = AAZStrArg(
+            options=["--service-name"],
+            arg_group="Parameters",
+            help="The name of the service.",
+            required=True,
+            enum={"AutomaticRepairs": "AutomaticRepairs"},
+        )
+        return cls._args_schema
+
+    def _execute_operations(self):
+        self.pre_operations()
+        yield self.VirtualMachineScaleSetsSetOrchestrationServiceState(ctx=self.ctx)()
+        self.post_operations()
+
+    @register_callback
+    def pre_operations(self):
+        pass
+
+    @register_callback
+    def post_operations(self):
+        pass
+
+    class VirtualMachineScaleSetsSetOrchestrationServiceState(AAZHttpOperation):
+        CLIENT_TYPE = "MgmtClient"
+
+        def __call__(self, *args, **kwargs):
+            request = self.make_request()
+            session = self.client.send_request(request=request, stream=False, **kwargs)
+            if session.http_response.status_code in [202]:
+                return self.client.build_lro_polling(
+                    self.ctx.args.no_wait,
+                    session,
+                    self.on_200,
+                    self.on_error,
+                    lro_options={"final-state-via": "location"},
+                    path_format_arguments=self.url_parameters,
+                )
+            if session.http_response.status_code in [200]:
+                return self.client.build_lro_polling(
+                    self.ctx.args.no_wait,
+                    session,
+                    self.on_200,
+                    self.on_error,
+                    lro_options={"final-state-via": "location"},
+                    path_format_arguments=self.url_parameters,
+                )
+
+            return self.on_error(session.http_response)
+
+        @property
+        def url(self):
+            return self.client.format_url(
+                "/subscriptions/{subscriptionId}/resourceGroups/{resourceGroupName}/providers/Microsoft.Compute/virtualMachineScaleSets/{vmScaleSetName}/setOrchestrationServiceState",
+                **self.url_parameters
+            )
+
+        @property
+        def method(self):
+            return "POST"
+
+        @property
+        def error_format(self):
+            return "ODataV4Format"
+
+        @property
+        def url_parameters(self):
+            parameters = {
+                **self.serialize_url_param(
+                    "resourceGroupName", self.ctx.args.resource_group,
+                    required=True,
+                ),
+                **self.serialize_url_param(
+                    "subscriptionId", self.ctx.subscription_id,
+                    required=True,
+                ),
+                **self.serialize_url_param(
+                    "vmScaleSetName", self.ctx.args.vm_scale_set_name,
+                    required=True,
+                ),
+            }
+            return parameters
+
+        @property
+        def query_parameters(self):
+            parameters = {
+                **self.serialize_query_param(
+                    "api-version", "2024-11-01",
+                    required=True,
+                ),
+            }
+            return parameters
+
+        @property
+        def header_parameters(self):
+            parameters = {
+                **self.serialize_header_param(
+                    "Content-Type", "application/json",
+                ),
+            }
+            return parameters
+
+        @property
+        def content(self):
+            _content_value, _builder = self.new_content_builder(
+                self.ctx.args,
+                typ=AAZObjectType,
+                typ_kwargs={"flags": {"required": True, "client_flatten": True}}
+            )
+            _builder.set_prop("action", AAZStrType, ".action", typ_kwargs={"flags": {"required": True}})
+            _builder.set_prop("serviceName", AAZStrType, ".service_name", typ_kwargs={"flags": {"required": True}})
+
+            return self.serialize_content(_content_value)
+
+        def on_200(self, session):
+            pass
+
+
+class _SetOrchestrationServiceStateHelper:
+    """Helper class for SetOrchestrationServiceState"""
+
+
+__all__ = ["SetOrchestrationServiceState"]

src/azure-cli/azure/cli/command_modules/vm/commands.py

@@ -404,6 +404,7 @@ def load_command_table(self, _):
         g.custom_show_command('identity show', 'show_vmss_identity')
         g.custom_command('deallocate', 'deallocate_vmss', supports_no_wait=True)
         g.custom_command('reimage', 'reimage_vmss', supports_no_wait=True)
+        g.custom_command('set-orchestration-service-state', 'set_orchestration_service_state', supports_no_wait=True)
 
     with self.command_group('vmss application', operation_group='virtual_machine_scale_sets') as g:
         g.custom_command('set', 'set_vmss_applications', validator=process_set_applications_namespace)
@@ -422,7 +423,6 @@ def load_command_table(self, _):
         g.generic_update_command('update', getter_name='get_vmss_modified_by_aaz', setter_name='update_vmss', supports_no_wait=True, command_type=compute_custom, validator=validate_vmss_update_namespace)
         g.custom_command('update-instances', 'update_vmss_instances', supports_no_wait=True)
         g.wait_command('wait', getter_name='get_vmss', getter_type=compute_custom)
-        g.custom_command('set-orchestration-service-state', 'set_orchestration_service_state', supports_no_wait=True)
 
         from .aaz.latest.vmss import List as VMSSList
         self.command_table['vmss list'] = VMSSList(loader=self,

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

@@ -5256,11 +5256,15 @@ def set_orchestration_service_state(cmd, resource_group_name, vm_scale_set_name,
     # currently service_name has only one available value "AutomaticRepairs". And SDK does not accept service_name,
     # instead SDK assign it to "AutomaticRepairs" in its own logic. As there may be more service name to be supported,
     # we define service_name as a required parameter here to avoid introducing a breaking change in the future.
-    client = _compute_client_factory(cmd.cli_ctx)
-    OrchestrationServiceStateInput = cmd.get_models('OrchestrationServiceStateInput')
-    state_input = OrchestrationServiceStateInput(service_name=service_name, action=action)
-    return sdk_no_wait(no_wait, client.virtual_machine_scale_sets.begin_set_orchestration_service_state,
-                       resource_group_name, vm_scale_set_name, state_input)
+    from .aaz.latest.vmss import SetOrchestrationServiceState as VmssSetOrchestrationServiceState
+    command_args = {
+        'resource_group': resource_group_name,
+        'vm_scale_set_name': vm_scale_set_name,
+        'action': action,
+        'service_name': service_name,
+        'no_wait': no_wait
+    }
+    return VmssSetOrchestrationServiceState(cli_ctx=cmd.cli_ctx)(command_args=command_args)
 
 
 def upgrade_vmss_extension(cmd, resource_group_name, vm_scale_set_name, no_wait=False):