From 1311b399258733bd72714cbd61faef2de518ed97 Mon Sep 17 00:00:00 2001
From: Tarun Gavara <tarungavara@gmail.com>
Date: Wed, 3 Jun 2020 10:59:18 +0530
Subject: [PATCH] SDK-407, JUPY-884: Add '--upload-to-source' param for
 JupyterNotebookCommand (#321)

---
 qds_sdk/commands.py   | 17 +++++++++++---
 tests/test_command.py | 52 +++++++++++++++++++++++++++++++++++++++++--
 2 files changed, 64 insertions(+), 5 deletions(-)

diff --git a/qds_sdk/commands.py b/qds_sdk/commands.py
index 99cc14ef..0966a9e2 100755
--- a/qds_sdk/commands.py
+++ b/qds_sdk/commands.py
@@ -1387,9 +1387,10 @@ class JupyterNotebookCommand(Command):
     optparser.add_option("--print-logs-live", action="store_true",
                          dest="print_logs_live", default=False, help="Fetch logs \
                          and print them to stderr while command is running.")
-    optparser.add_option("--skip-upload-to-source", action="store_false",
-                         dest="upload_to_source", default=True, help="Do not \
-                         upload notebook to source after completion of execution")
+    optparser.add_option("--upload-to-source", dest="upload_to_source", default='true',
+                         help="Upload notebook to source after completion of \
+                         execution. Specify the value as either 'true' or 'false'.\
+                         Default value is 'true'.")
 
     @classmethod
     def parse(cls, args):
@@ -1417,6 +1418,16 @@ def parse(cls, args):
                 options.macros = validate_json_input(options.macros, 'Macros', cls)
             if options.retry is not None:
                 options.retry = int(options.retry)
+            if options.upload_to_source is not None:
+                options.upload_to_source = options.upload_to_source.lower()
+                if options.upload_to_source == 'true':
+                    options.upload_to_source = True
+                elif options.upload_to_source == 'false':
+                    options.upload_to_source = False
+                else:
+                    msg = "Upload to Source parameter takes a value of either 'true' \
+                    or 'false' only."
+                    raise ParseError(msg, cls.optparser.format_help())
         except OptionParsingError as e:
             raise ParseError(e.msg, cls.optparser.format_help())
         except OptionParsingExit as e:
diff --git a/tests/test_command.py b/tests/test_command.py
index bbbf83a6..84c8955a 100644
--- a/tests/test_command.py
+++ b/tests/test_command.py
@@ -2247,9 +2247,50 @@ def test_submit_pool(self):
                  'can_notify': False,
                  'pool': 'batch'})
 
-    def test_submit_skip_upload_to_source(self):
+    def test_submit_no_upload_to_source(self):
+        sys.argv = ['qds.py', 'jupyternotebookcmd', 'submit', '--path', 'folder/file']
+        print_command()
+        Connection._api_call = Mock(return_value={'id': 1234})
+        qds.main()
+        Connection._api_call.assert_called_with('POST', 'commands',
+                {'retry': None,
+                 'name': None,
+                 'tags': None,
+                 'label': None,
+                 'macros': None,
+                 'arguments': None,
+                 'timeout': None,
+                 'path': 'folder/file',
+                 'retry_delay': None,
+                 'command_type': 'JupyterNotebookCommand',
+                 'upload_to_source': True,
+                 'can_notify': False,
+                 'pool': None})
+
+    def test_submit_upload_to_source(self):
+        sys.argv = ['qds.py', 'jupyternotebookcmd', 'submit', '--path', 'folder/file',
+                    '--upload-to-source', 'True']
+        print_command()
+        Connection._api_call = Mock(return_value={'id': 1234})
+        qds.main()
+        Connection._api_call.assert_called_with('POST', 'commands',
+                {'retry': None,
+                 'name': None,
+                 'tags': None,
+                 'label': None,
+                 'macros': None,
+                 'arguments': None,
+                 'timeout': None,
+                 'path': 'folder/file',
+                 'retry_delay': None,
+                 'command_type': 'JupyterNotebookCommand',
+                 'upload_to_source': True,
+                 'can_notify': False,
+                 'pool': None})
+
+    def test_submit_upload_to_source_false(self):
         sys.argv = ['qds.py', 'jupyternotebookcmd', 'submit', '--path', 'folder/file',
-                    '--skip-upload-to-source']
+                    '--upload-to-source', 'False']
         print_command()
         Connection._api_call = Mock(return_value={'id': 1234})
         qds.main()
@@ -2268,6 +2309,13 @@ def test_submit_skip_upload_to_source(self):
                  'can_notify': False,
                  'pool': None})
 
+    def test_submit_upload_to_source_wrong_param(self):
+        sys.argv = ['qds.py', 'jupyternotebookcmd', 'submit', '--path', 'folder/file',
+                    '--upload-to-source', 'wrong']
+        print_command()
+        with self.assertRaises(qds_sdk.exception.ParseError):
+            qds.main()
+
     def test_submit_retry(self):
         sys.argv = ['qds.py', 'jupyternotebookcmd', 'submit', '--path', 'folder/file',
                     '--retry', '1']