SDK-416 : Add managed disk type for azure cluster (#330)

(cherry picked from commit fc18ceb)
qubole · Oct 6, 2020 · 29834ce · 29834ce
1 parent 5f1ab81
commit 29834ce
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 2 deletions.
diff --git a/qds_sdk/cloud/azure_cloud.py b/qds_sdk/cloud/azure_cloud.py
@@ -22,6 +22,7 @@ def set_cloud_config(self,
                          storage_account_name=None,
                          disk_storage_account_name=None,
                          disk_storage_account_resource_group_name=None,
+                         managed_disk_account_type=None,
                          persistent_security_groups=None,
                          bastion_node_public_dns=None,
                          vnet_name=None,
@@ -56,6 +57,8 @@ def set_cloud_config(self,
             disk_storage_account_resource_group_name: Disk storage account resource group
                 namefor azure cluster
 
+            managed_disk_account_type:  Managed Disk Type for azure cluster
+
             persistent_security_groups: security group to associate with each node of the cluster.
                 Typically used to provide access to external hosts
 
@@ -87,7 +90,8 @@ def set_cloud_config(self,
                                 vnet_resource_group_name, master_static_nic_name, master_static_public_ip_name)
         self.set_storage_config(storage_access_key, storage_account_name,
                                 disk_storage_account_name,
-                                disk_storage_account_resource_group_name)
+                                disk_storage_account_resource_group_name,
+                                managed_disk_account_type)
         self.resource_group_name = resource_group_name
 
     def set_compute_config(self,
@@ -128,12 +132,14 @@ def set_storage_config(self,
                            storage_access_key=None,
                            storage_account_name=None,
                            disk_storage_account_name=None,
-                           disk_storage_account_resource_group_name=None):
+                           disk_storage_account_resource_group_name=None,
+                           managed_disk_account_type=None):
         self.storage_config['storage_access_key'] = storage_access_key
         self.storage_config['storage_account_name'] = storage_account_name
         self.storage_config['disk_storage_account_name'] = disk_storage_account_name
         self.storage_config['disk_storage_account_resource_group_name'] \
             = disk_storage_account_resource_group_name
+        self.storage_config['managed_disk_account_type'] = managed_disk_account_type
 
     def set_cloud_config_from_arguments(self, arguments):
         self.set_cloud_config(compute_client_id=arguments.compute_client_id,
@@ -146,6 +152,7 @@ def set_cloud_config_from_arguments(self, arguments):
                               storage_account_name=arguments.storage_account_name,
                               disk_storage_account_name=arguments.disk_storage_account_name,
                               disk_storage_account_resource_group_name=arguments.disk_storage_account_resource_group_name,
+                              managed_disk_account_type=arguments.managed_disk_account_type,
                               vnet_name=arguments.vnet_name,
                               subnet_name=arguments.subnet_name,
                               vnet_resource_group_name=arguments.vnet_resource_group_name,
@@ -312,3 +319,7 @@ def create_parser(self, argparser):
                                     dest="disk_storage_account_resource_group_name",
                                     default=None,
                                     help="disk storage account resource group for azure cluster")
+        storage_config.add_argument("--managed-disk-account-type",
+                                    dest="managed_disk_account_type",
+                                    default=None,
+                                    help="managed disk type for azure cluster")
diff --git a/tests/test_clusterv2.py b/tests/test_clusterv2.py
@@ -219,6 +219,26 @@ def test_azure_storage_config(self):
                                                       }
                                                  })
 
+    def test_azure_managed_disk_config(self):
+        sys.argv = ['qds.py', '--version', 'v2', '--cloud', 'AZURE', 'cluster', 'create', '--label', 'test_label',
+                    '--storage-access-key', 'testkey', '--storage-account-name', 'test_account_name',
+                    '--managed-disk-account-type', 'test_managed_disk']
+        Qubole.cloud = None
+        print_command()
+        Connection._api_call = Mock(return_value={})
+        qds.main()
+        Connection._api_call.assert_called_with('POST', 'clusters',
+                                                {'cluster_info':
+                                                     {'label': ['test_label']},
+                                                 'cloud_config':
+                                                     {'storage_config':
+                                                          {'storage_access_key': 'testkey',
+                                                           'storage_account_name': 'test_account_name',
+                                                           'managed_disk_account_type': 'test_managed_disk'
+                                                           }
+                                                      }
+                                                 })
+
     def test_azure_network_config(self):
         sys.argv = ['qds.py', '--version', 'v2', '--cloud', 'AZURE', 'cluster', 'create', '--label', 'test_label',
                     '--vnet-name', 'testvnet', '--subnet-name', 'testsubnet',