From 3dc5bc426a0370fe08fcea1fde697faa7ec1628e Mon Sep 17 00:00:00 2001
From: saiyam1712 <56955417+saiyam1712@users.noreply.github.com>
Date: Mon, 4 Nov 2019 11:37:00 +0530
Subject: [PATCH] SDK-381 support for spark streaming cluster (#289)

support for spark streaming cluster
---
 qds_sdk/engine.py       |  2 +-
 tests/test_clusterv2.py | 19 ++++++++++++++++++-
 2 files changed, 19 insertions(+), 2 deletions(-)

diff --git a/qds_sdk/engine.py b/qds_sdk/engine.py
index 3070f699..ee55e50d 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", "hbase", "airflow", "deeplearning"],
+                                  choices=["hadoop", "hadoop2", "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 eec62ade..4e5dd0ff 100644
--- a/tests/test_clusterv2.py
+++ b/tests/test_clusterv2.py
@@ -460,6 +460,23 @@ def test_spark_engine_config(self):
                                                               'custom_spark_config': 'spark-overrides'}},
                                                      'cluster_info': {'label': ['test_label'],}})
 
+    def test_sparkstreaming_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', 'sparkstreaming', '--custom-spark-config', 'spark-overrides']
+            Qubole.cloud = None
+            print_command()
+            Connection._api_call = Mock(return_value={})
+            qds.main()
+            Connection._api_call.assert_called_with('POST', 'clusters',
+                                                {'engine_config':
+                                                     {'flavour': 'sparkstreaming',
+                                                      'spark_settings': {
+                                                          'custom_spark_config': 'spark-overrides'}},
+                                                 'cluster_info': {'label': ['test_label'],}})
+
     def test_airflow_engine_config(self):
         with tempfile.NamedTemporaryFile() as temp:
             temp.write("config.properties:\na=1\nb=2".encode("utf8"))
@@ -787,7 +804,7 @@ def test_engine_config(self):
             temp.write("a=1\nb=2".encode("utf8"))
             temp.flush()
             sys.argv = ['qds.py', '--version', 'v2', 'cluster', 'update', '123',
-                        '--use-qubole-placement-policy', '--enable-rubix', 
+                        '--use-qubole-placement-policy', '--enable-rubix',
                         '--custom-hadoop-config',temp.name]
             Qubole.cloud = None
             print_command()