From 6ce017b8cad3f672f661bf3855f27c05db852c0b Mon Sep 17 00:00:00 2001
From: satyabolnedi <57746175+satyabolnedi@users.noreply.github.com>
Date: Fri, 29 Nov 2019 13:29:31 +0530
Subject: [PATCH] ACM-4728: Add HS2 cluster creation from qds-sdk (#290)

Added support for creating HS2 cluster
---
 qds_sdk/clusterv2.py    | 11 +++++++++--
 qds_sdk/engine.py       |  2 +-
 tests/test_clusterv2.py | 19 +++++++++++++++++++
 3 files changed, 29 insertions(+), 3 deletions(-)

diff --git a/qds_sdk/clusterv2.py b/qds_sdk/clusterv2.py
index e65d6570..e0e4cb0c 100755
--- a/qds_sdk/clusterv2.py
+++ b/qds_sdk/clusterv2.py
@@ -181,7 +181,8 @@ def set_cluster_info_from_arguments(self, arguments):
                               disable_cluster_pause=arguments.disable_cluster_pause,
                               paused_cluster_timeout_mins=arguments.paused_cluster_timeout_mins,
                               disable_autoscale_node_pause=arguments.disable_autoscale_node_pause,
-                              paused_autoscale_node_timeout_mins=arguments.paused_autoscale_node_timeout_mins)
+                              paused_autoscale_node_timeout_mins=arguments.paused_autoscale_node_timeout_mins,
+                              parent_cluster_id=arguments.parent_cluster_id)
 
     def set_cluster_info(self,
                          disallow_cluster_termination=None,
@@ -223,7 +224,8 @@ def set_cluster_info(self,
                          disable_cluster_pause=None,
                          paused_cluster_timeout_mins=None,
                          disable_autoscale_node_pause=None,
-                         paused_autoscale_node_timeout_mins=None):
+                         paused_autoscale_node_timeout_mins=None,
+                         parent_cluster_id=None):
         """
         Args:
 
@@ -357,6 +359,7 @@ def set_cluster_info(self,
 
         self.cluster_info['rootdisk'] = {}
         self.cluster_info['rootdisk']['size'] = root_disk_size
+        self.cluster_info['parent_cluster_id'] = parent_cluster_id
 
         self.set_spot_instance_settings(maximum_bid_price_percentage, timeout_for_request,
                                         maximum_spot_instance_percentage)
@@ -515,6 +518,10 @@ def cluster_info_parser(argparser, action):
                                   dest="root_disk_size",
                                   type=int,
                                   help="size of the root volume in GB")
+        cluster_info.add_argument("--parent-cluster-id",
+                                  dest="parent_cluster_id",
+                                  type=int,
+                                  help="Id of the parent cluster this hs2 cluster is attached to")
         termination = cluster_info.add_mutually_exclusive_group()
         termination.add_argument("--disallow-cluster-termination",
                                  dest="disallow_cluster_termination",
diff --git a/qds_sdk/engine.py b/qds_sdk/engine.py
index ee55e50d..22893df8 100644
--- a/qds_sdk/engine.py
+++ b/qds_sdk/engine.py
@@ -147,7 +147,7 @@ def engine_parser(argparser):
         engine_group = argparser.add_argument_group("engine settings")
         engine_group.add_argument("--flavour",
                                   dest="flavour",
-                                  choices=["hadoop", "hadoop2", "presto", "spark", "sparkstreaming", "hbase", "airflow", "deeplearning"],
+                                  choices=["hadoop", "hadoop2", "hs2", "presto", "spark", "sparkstreaming", "hbase", "airflow", "deeplearning"],
                                   default=None,
                                   help="Set engine flavour")
 
diff --git a/tests/test_clusterv2.py b/tests/test_clusterv2.py
index 34dd8433..9acc5dc5 100644
--- a/tests/test_clusterv2.py
+++ b/tests/test_clusterv2.py
@@ -443,6 +443,25 @@ def test_presto_engine_config(self):
                                                           }},
                                                      'cluster_info': {'label': ['test_label']}})
 
+    def test_hs2_engine_config(self):
+        with tempfile.NamedTemporaryFile() as temp:
+            temp.write("config.properties:\na=1\nb=2".encode("utf8"))
+            temp.flush()
+            sys.argv = ['qds.py', '--version', 'v2', 'cluster', 'create', '--label',
+                         'test_label', '--flavour', 'hs2', '--node-bootstrap-file', 'test_file_name', '--slave-instance-type', 'c1.xlarge', '--min-nodes', '3', '--parent-cluster-id', '1']
+            Qubole.cloud = None
+            print_command()
+            Connection._api_call = Mock(return_value={})
+            qds.main()
+            Connection._api_call.assert_called_with('POST', 'clusters',
+                                                    {'engine_config':
+                                                         {'flavour': 'hs2'},
+                                                    'cluster_info': {'label': ['test_label'],
+                                                    'parent_cluster_id': 1, 
+                                                    'min_nodes': 3,
+                                                    'node_bootstrap': 'test_file_name',
+                                                    'slave_instance_type': 'c1.xlarge' }})
+
     def test_spark_engine_config(self):
         with tempfile.NamedTemporaryFile() as temp:
             temp.write("config.properties:\na=1\nb=2".encode("utf8"))