From 3694711a7e975324d52c258ab73a8f5e766a3f1c Mon Sep 17 00:00:00 2001
From: Fred-sun <37327967+Fred-sun@users.noreply.github.com>
Date: Wed, 3 Apr 2019 11:08:27 +0800
Subject: [PATCH] Update azure_rm_managed_disk.py --add zones (#53788)

---
 .../cloud/azure/azure_rm_managed_disk.py      | 24 +++++++++++++++++--
 .../azure_rm_managed_disk/tasks/main.yml      | 20 ++++++++++++++++
 2 files changed, 42 insertions(+), 2 deletions(-)

diff --git a/lib/ansible/modules/cloud/azure/azure_rm_managed_disk.py b/lib/ansible/modules/cloud/azure/azure_rm_managed_disk.py
index 1603b3e8281..f5a0f05b1cd 100644
--- a/lib/ansible/modules/cloud/azure/azure_rm_managed_disk.py
+++ b/lib/ansible/modules/cloud/azure/azure_rm_managed_disk.py
@@ -93,6 +93,15 @@ options:
     tags:
         description:
             - Tags to assign to the managed disk.
+    zone:
+        description:
+            - "Allowed values: 1, 2, 3, ''."
+        choices:
+            - 1
+            - 2
+            - 3
+            - ''
+        version_added: "2.8"
 
 extends_documentation_fragment:
     - azure
@@ -182,7 +191,8 @@ def managed_disk_to_dict(managed_disk):
         disk_size_gb=managed_disk.disk_size_gb,
         os_type=managed_disk.os_type.lower() if managed_disk.os_type else None,
         storage_account_type=managed_disk.sku.name if managed_disk.sku else None,
-        managed_by=managed_disk.managed_by
+        managed_by=managed_disk.managed_by,
+        zone=managed_disk.zones[0] if managed_disk.zones and len(managed_disk.zones) > 0 else ''
     )
 
 
@@ -228,6 +238,10 @@ class AzureRMManagedDisk(AzureRMModuleBase):
             ),
             managed_by=dict(
                 type='str'
+            ),
+            zone=dict(
+                type='str',
+                choices=['', '1', '2', '3']
             )
         )
         required_if = [
@@ -248,6 +262,7 @@ class AzureRMManagedDisk(AzureRMModuleBase):
         self.os_type = None
         self.disk_size_gb = None
         self.tags = None
+        self.zone = None
         self.managed_by = None
         super(AzureRMManagedDisk, self).__init__(
             derived_arg_spec=self.module_arg_spec,
@@ -338,11 +353,13 @@ class AzureRMManagedDisk(AzureRMModuleBase):
             self.fail("Error getting virtual machine {0} - {1}".format(name, str(exc)))
 
     def generate_managed_disk_property(self):
-        # TODO: Add support for EncryptionSettings, DiskIOPSReadWrite, DiskMBpsReadWrite, Zones
+        # TODO: Add support for EncryptionSettings, DiskIOPSReadWrite, DiskMBpsReadWrite
         disk_params = {}
         creation_data = {}
         disk_params['location'] = self.location
         disk_params['tags'] = self.tags
+        if self.zone:
+            disk_params['zones'] = [self.zone]
         if self.storage_account_type:
             storage_account_type = self.compute_models.DiskSku(name=self.storage_account_type)
             disk_params['sku'] = storage_account_type
@@ -393,6 +410,9 @@ class AzureRMManagedDisk(AzureRMModuleBase):
         if new_disk.get('tags') is not None:
             if not found_disk['tags'] == new_disk['tags']:
                 resp = True
+        if self.zone is not None:
+            if not found_disk['zone'] == self.zone:
+                resp = True
         return resp
 
     def delete_managed_disk(self):
diff --git a/test/integration/targets/azure_rm_managed_disk/tasks/main.yml b/test/integration/targets/azure_rm_managed_disk/tasks/main.yml
index e245f15344c..6f93dd10b5b 100644
--- a/test/integration/targets/azure_rm_managed_disk/tasks/main.yml
+++ b/test/integration/targets/azure_rm_managed_disk/tasks/main.yml
@@ -11,6 +11,7 @@
        state: absent
    with_items:
      - "md{{ rpfx }}os"
+     - "md{{ rpfx }}3"
      - "md{{ rpfx }}2"
      - "md{{ rpfx }}1"
 
@@ -64,6 +65,24 @@
        - output.state.os_type == None
        - output.state.storage_account_type == "Standard_LRS"
 
+ - name: Create new managed disk with zone
+   azure_rm_managed_disk:
+       resource_group: "{{ resource_group }}"
+       name: "md{{ rpfx }}3"
+       storage_account_type: "Standard_LRS"
+       disk_size_gb: 1
+       zone: 1
+       tags:
+           testing: testing
+           delete: never
+   register: output
+
+ - name: Assert status succeeded and results include an zone value
+   assert:
+     that:
+       - output.changed
+       - output.state.zone == ["1"] 
+
  - name: Change the operating system type of the managed disk to linux
    azure_rm_managed_disk:
        resource_group: "{{ resource_group }}"
@@ -487,6 +506,7 @@
    with_items:
        - 1
        - 2
+       - 3
 
  - name: Delete virtual machine
    azure_rm_virtualmachine: