Make exporting to directory munge target_path (#20)

Author: andrewmchen

databricks_cli/configure/config.py

@@ -39,6 +39,9 @@
 
 def require_config(function):
     def decorator(*args, **kwargs):
+        if kwargs.get('test_mode'):
+            del kwargs['test_mode']
+            return function(*args, **kwargs)
         config = DatabricksConfig.fetch_from_fs()
         if not config.is_valid:
             error_and_quit(('You haven\'t configured the CLI yet! '

databricks_cli/workspace/api.py

@@ -110,8 +110,12 @@ def import_workspace(source_path, target_path, language, fmt, is_overwrite):
 
 
 def export_workspace(source_path, target_path, fmt, is_overwrite):
+    """
+    Faithfully exports the source_path to the target_path. Does not
+    attempt to do any munging of the target_path if it is a directory.
+    """
     if os.path.exists(target_path) and not is_overwrite:
-        raise LocalFileExistsException()
+        raise LocalFileExistsException('Target {} already exists.'.format(target_path))
     workspace_client = get_workspace_client()
     output = workspace_client.export_workspace(source_path, fmt)
     content = output['content']

databricks_cli/workspace/cli.py

@@ -110,6 +110,12 @@ def export_workspace_cli(source_path, target_path, format, overwrite): # NOQA
     format is documented at
     https://docs.databricks.com/api/latest/workspace.html#notebookexportformat.
     """
+    if os.path.isdir(target_path):
+        file_info = get_status(source_path)
+        if not file_info.is_notebook:
+            raise RuntimeError('Export can only be called on a notebook.')
+        extension = WorkspaceLanguage.to_extension(file_info.language)
+        target_path = os.path.join(target_path, file_info.basename + extension)
     export_workspace(source_path, target_path, format, overwrite) # NOQA
 
 

tests/configure/test_config.py

@@ -40,6 +40,14 @@ def test_function(x):
         assert test_function(1) == 1
 
 
+def test_require_config_valid_test_mode():
+    @config.require_config
+    def test_function(x):
+        return x
+
+    assert test_function(1, test_mode=True) == 1 # noqa
+
+
 def test_require_config_invalid():
     with mock.patch('databricks_cli.configure.config.DatabricksConfig') as DatabricksConfigMock:
         with mock.patch('databricks_cli.configure.config.error_and_quit') as error_and_quit_mock:

tests/workspace/test_cli.py

@@ -26,8 +26,16 @@
 
 import databricks_cli.workspace.cli as cli
 import databricks_cli.workspace.api as api
-from databricks_cli.workspace.api import WorkspaceFileInfo
-from databricks_cli.workspace.types import WorkspaceLanguage
+from databricks_cli.workspace.api import WorkspaceFileInfo, NOTEBOOK
+from databricks_cli.workspace.types import WorkspaceLanguage, WorkspaceFormat
+
+def test_export_workspace_cli(tmpdir):
+    path = tmpdir.strpath
+    with mock.patch('databricks_cli.workspace.cli.get_status') as get_status_mock:
+        with mock.patch('databricks_cli.workspace.cli.export_workspace') as export_workspace_mock:
+            get_status_mock.return_value = WorkspaceFileInfo('/notebook-name', NOTEBOOK, WorkspaceLanguage.SCALA)
+            cli.export_workspace_cli.callback('/notebook-name', path, WorkspaceFormat.SOURCE, False, test_mode=True)
+            assert export_workspace_mock.call_args[0][1] == os.path.join(path, 'notebook-name.scala')
 
 def test_export_dir_helper(tmpdir):
     """