Add non-premium disk option for repair VM

src/vm-repair/HISTORY.rst

@@ -2,10 +2,15 @@
 Release History
 ===============
 
+2.0.4
+++++++
+Added new parameter `os-disk-type` to `vm repair create` to let users specify the repair vm's os disk storage account type. 
+
 2.0.3
 ++++++
 Added new long parameter functionality in `vm repair run` cmd `parameters` parameter. When using the prefix `++`, the entire key=value string will be sent to the running script, not just the value.
 
+
 2.0.2
 ++++++
 Updated parameter descriptions and examples for `az vm repair create`.

src/vm-repair/azext_vm_repair/_help.py

@@ -37,6 +37,9 @@
         - name: Create a repair VM from a source VM with an encrypted disk. The repair VM is created with the data disk unencrypted and accessible.
           text: >
             az vm repair create -g MyResourceGroup -n myVM --yes --repair-username <username> --repair-password <password> --unlock-encrypted-vm --encrypt-recovery-key <key>
+        - name: Create a repair VM with an OS Disk storage type of StandardSSD_LRS.
+          text: >
+            az vm repair create -g MyResourceGroup -n myVM --yes --repair-username <username> --repair-password <password> --os-disk-type StandardSSD_LRS
 """
 
 helps['vm repair restore'] = """

src/vm-repair/azext_vm_repair/_params.py

@@ -35,6 +35,7 @@ def load_arguments(self, _):
         c.argument('distro', help='Option to create repair vm from a specific linux distro (rhel7|rhel8|sles12|sles15|ubuntu20|centos7|centos8|oracle7)')
         c.argument('yes', help='Option to skip prompt for associating public ip in no Tty mode')
         c.argument('disable_trusted_launch', help='Option to disable Trusted Launch security type on the repair vm by setting the security type to Standard.')
+        c.argument('os_disk_type', help='Change the OS Disk storage type from the default of PremiumSSD_LRS.')
 
     with self.argument_context('vm repair restore') as c:
         c.argument('repair_vm_id', help='Repair VM resource id.')

src/vm-repair/azext_vm_repair/custom.py

@@ -56,7 +56,7 @@
 logger = get_logger(__name__)
 
 
-def create(cmd, vm_name, resource_group_name, repair_password=None, repair_username=None, repair_vm_name=None, copy_disk_name=None, repair_group_name=None, unlock_encrypted_vm=False, enable_nested=False, associate_public_ip=False, distro='ubuntu', yes=False, encrypt_recovery_key="", disable_trusted_launch=False):  
+def create(cmd, vm_name, resource_group_name, repair_password=None, repair_username=None, repair_vm_name=None, copy_disk_name=None, repair_group_name=None, unlock_encrypted_vm=False, enable_nested=False, associate_public_ip=False, distro='ubuntu', yes=False, encrypt_recovery_key="", disable_trusted_launch=False, os_disk_type=None):  
     """  
     This function creates a repair VM.  
 
@@ -75,19 +75,17 @@ def create(cmd, vm_name, resource_group_name, repair_password=None, repair_usern
     - distro: The Linux distribution to use for the repair VM. Default is 'ubuntu'.  
     - yes: If True, confirmation prompts will be skipped. Default is False.  
     - encrypt_recovery_key: The Bitlocker recovery key to use for encrypting the VM. Default is an empty string.  
-    - disable_trusted_launch: A flag parameter that, when used, sets the security type of the repair VM to Standard. 
+    - disable_trusted_launch: A flag parameter that, when used, sets the security type of the repair VM to Standard.
+    - os_disk_type: A flag parameter that, when used, sets the OS disk type of the repair VM to the specified type. 
     """  
 
-    # A warning about a future breaking change.  
-    logger.warning('After the November 2024 release, if the image of the source Windows VM is not found, the \'az vm repair create\' command will default to use a 2022-Datacenter image for the repair VM.')  
-
     # Logging all the command parameters, except the sensitive data.
     # Mask sensitive information
     masked_repair_password = '****' if repair_password else None
     masked_repair_username = '****' if repair_username else None
     masked_repair_encrypt_recovery_key = '****' if encrypt_recovery_key else None
-    logger.debug('vm repair create command parameters: vm_name: %s, resource_group_name: %s, repair_password: %s, repair_username: %s, repair_vm_name: %s, copy_disk_name: %s, repair_group_name: %s, unlock_encrypted_vm: %s, enable_nested: %s, associate_public_ip: %s, distro: %s, yes: %s, encrypt_recovery_key: %s, disable_trusted_launch: %s', 
-                 vm_name, resource_group_name, masked_repair_password, masked_repair_username, repair_vm_name, copy_disk_name, repair_group_name, unlock_encrypted_vm, enable_nested, associate_public_ip, distro, yes, masked_repair_encrypt_recovery_key, disable_trusted_launch)  
+    logger.debug('vm repair create command parameters: vm_name: %s, resource_group_name: %s, repair_password: %s, repair_username: %s, repair_vm_name: %s, copy_disk_name: %s, repair_group_name: %s, unlock_encrypted_vm: %s, enable_nested: %s, associate_public_ip: %s, distro: %s, yes: %s, encrypt_recovery_key: %s, disable_trusted_launch: %s, os_disk_type: %s', 
+                 vm_name, resource_group_name, masked_repair_password, masked_repair_username, repair_vm_name, copy_disk_name, repair_group_name, unlock_encrypted_vm, enable_nested, associate_public_ip, distro, yes, masked_repair_encrypt_recovery_key, disable_trusted_launch, os_disk_type)  
 
     # Initializing a command helper object.  
     command = command_helper(logger, cmd, 'vm repair create')  
@@ -205,8 +203,11 @@ def create(cmd, vm_name, resource_group_name, repair_password=None, repair_usern
             logger.info('Creating resource group for repair VM and its resources...')  
             _call_az_command(create_resource_group_command)  
 
+        # Check if user is changing the Repair VM os disk type
+        if os_disk_type:
+            create_repair_vm_command += ' --storage-sku {os_disk_type} '.format(os_disk_type=os_disk_type)
 
-                # Check if the source VM uses managed disks.  
+        # Check if the source VM uses managed disks.  
         # If it does, the repair VM will also be created with managed disks.  
         if is_managed:  
             logger.info('Source VM uses managed disks. Creating repair VM with managed disks.\n')  
@@ -218,7 +219,6 @@ def create(cmd, vm_name, resource_group_name, repair_password=None, repair_usern
             # The command includes the resource group name, copy disk name, target disk name, SKU, location, and OS type.  
             copy_disk_command = 'az disk create -g {g} -n {n} --source {s} --sku {sku} --location {loc} --os-type {os_type} --query id -o tsv' \
                 .format(g=resource_group_name, n=copy_disk_name, s=target_disk_name, sku=disk_sku, loc=location, os_type=os_type)  
-
             # If the Hyper-V generation for the disk is available, append it to the copy disk command.  
             if hyperV_generation:  
                 copy_disk_command += ' --hyper-v-generation {hyperV}'.format(hyperV=hyperV_generation)  
@@ -233,12 +233,11 @@ def create(cmd, vm_name, resource_group_name, repair_password=None, repair_usern
             if source_vm.zones:  
                 zone = source_vm.zones[0]  
                 copy_disk_command += ' --zone {zone}'.format(zone=zone)  
-
+    
             # Execute the command to create a copy of the OS disk of the source VM.  
             logger.info('Copying OS disk of source VM...')  
             copy_disk_id = _call_az_command(copy_disk_command).strip('\n')  
 
-
             # Depending on the operating system of the source VM and whether it's encrypted, different steps are taken.  
             # If the source VM is not a Linux machine, create the repair VM.  
             # This is the case for Windows VMs, both encrypted and not encrypted.  

src/vm-repair/azext_vm_repair/tests/latest/test_repair_commands.py

@@ -1180,7 +1180,9 @@ def test_vmrepair_WinRunRepairVMIdafterRefactor(self, resource_group):
 
         # Call Restore
         self.cmd('vm repair restore -g {rg} -n {vm} --yes')
+
 
+
 
 @pytest.mark.linuxLongParams
 class LinuxLongParametersParsing(LiveScenarioTest):
@@ -1242,5 +1244,65 @@ def test_vmrepair_LongoptionsParams(self, resource_group):
         log_contains_initiator_selfhelp = "initiator=selfhelp" in result_run['logs'].lower()  
         assert log_contains_initiator_selfhelp, "The logs do not contain 'initiator=selfhelp'"  
 
+        # Call Restore
+        self.cmd('vm repair restore -g {rg} -n {vm} --yes')
+
+
+@pytest.mark.NonPremiumOSDisk
+class WindowsNonPremiumOSDiskRepairVM(LiveScenarioTest):
+
+    @ResourceGroupPreparer(location='westus2')
+    def test_vmrepair_WinNonPremiumOSDiskRepairVM(self, resource_group):
+        import uuid
+        import secrets
+        import string
+        base_password = "Passw0rd2024"
+        guid_suffix = str(uuid.uuid4())
+        secure_password = base_password + guid_suffix
+        username_length = 8
+        secure_username = ''.join(secrets.choice(string.ascii_letters + string.digits) for _ in range(username_length))
+        stnd_os_disk_type = "StandardSSD_LRS"
+        self.kwargs.update({
+            'vm': 'vm1',
+            'admin_password': secure_password,
+            'admin_username': secure_username,
+            'stnd_os_disk_type': stnd_os_disk_type
+        })
+
+        # Create test VM
+        self.cmd('vm create -g {rg} -n {vm} --admin-username {admin_username} --admin-password {admin_password} --image MicrosoftWindowsServer:WindowsServer:2022-datacenter-g2:latest')
+        vms = self.cmd('vm list -g {rg} -o json').get_output_in_json()
+        # Something wrong with vm create command if it fails here
+        assert len(vms) == 1
+
+        # Create Repair VM
+        repair_vm = self.cmd('vm repair create -g {rg} -n {vm} --repair-username {admin_username} --repair-password {admin_password} --yes --os-disk-type {stnd_os_disk_type} -o json').get_output_in_json()
+        assert repair_vm['status'] == STATUS_SUCCESS, repair_vm['error_message']
+        # Check repair VM
+        self.kwargs.update({
+            'vm': 'vm1',
+            'admin_password': secure_password,
+            'admin_username': secure_username,
+            'repair_resource_group': repair_vm['repair_resource_group']
+        })
+        repair_vms = self.cmd('vm list -g {repair_resource_group} -o json').get_output_in_json()
+        assert len(repair_vms) == 1
+        repair_vm = repair_vms[0]
+        repair_vm_id = repair_vm['id']
+        # Verify the image sku is the default sku
+        os_disk_type = repair_vm['storageProfile']['osDisk']['managedDisk']['storageAccountType']
+        assert os_disk_type == stnd_os_disk_type, "Os Disk Storage Account Type is not the expected value."
+
+        repair_vm_id = repair_vm['id']
+        self.kwargs.update({
+            'vm': 'vm1',
+            'admin_password': secure_password,
+            'admin_username': secure_username,
+            'repair_vm_id': repair_vm_id
+        })
+        # Run a script for testing repair-vm-id
+        result_run = self.cmd('vm repair run -g {rg} -n {vm} --run-id win-hello-world --repair-vm-id {repair_vm_id} --run-on-repair -o json').get_output_in_json()
+        assert result_run['status'] == STATUS_SUCCESS, result_run['error_message']
+
         # Call Restore
         self.cmd('vm repair restore -g {rg} -n {vm} --yes')

src/vm-repair/setup.py

@@ -8,7 +8,7 @@
 from codecs import open
 from setuptools import setup, find_packages
 
-VERSION = "2.0.3"
+VERSION = "2.0.4"
 
 CLASSIFIERS = [
     'Development Status :: 4 - Beta',