github.com/databricks/cli@v0.203.0/internal/workspace_test.go (about)

     1  package internal
     2  
     3  import (
     4  	"context"
     5  	"errors"
     6  	"io"
     7  	"net/http"
     8  	"os"
     9  	"path"
    10  	"path/filepath"
    11  	"strings"
    12  	"testing"
    13  
    14  	"github.com/databricks/cli/libs/filer"
    15  	"github.com/databricks/databricks-sdk-go"
    16  	"github.com/databricks/databricks-sdk-go/apierr"
    17  	"github.com/stretchr/testify/assert"
    18  	"github.com/stretchr/testify/require"
    19  )
    20  
    21  func TestAccWorkspaceList(t *testing.T) {
    22  	t.Log(GetEnvOrSkipTest(t, "CLOUD_ENV"))
    23  
    24  	stdout, stderr := RequireSuccessfulRun(t, "workspace", "list", "/")
    25  	outStr := stdout.String()
    26  	assert.Contains(t, outStr, "ID")
    27  	assert.Contains(t, outStr, "Type")
    28  	assert.Contains(t, outStr, "Language")
    29  	assert.Contains(t, outStr, "Path")
    30  	assert.Equal(t, "", stderr.String())
    31  }
    32  
    33  func TestWorkpaceListErrorWhenNoArguments(t *testing.T) {
    34  	_, _, err := RequireErrorRun(t, "workspace", "list")
    35  	assert.Equal(t, "accepts 1 arg(s), received 0", err.Error())
    36  }
    37  
    38  func TestWorkpaceGetStatusErrorWhenNoArguments(t *testing.T) {
    39  	_, _, err := RequireErrorRun(t, "workspace", "get-status")
    40  	assert.Equal(t, "accepts 1 arg(s), received 0", err.Error())
    41  }
    42  
    43  func TestWorkpaceExportPrintsContents(t *testing.T) {
    44  	t.Log(GetEnvOrSkipTest(t, "CLOUD_ENV"))
    45  
    46  	ctx := context.Background()
    47  	w := databricks.Must(databricks.NewWorkspaceClient())
    48  	tmpdir := temporaryWorkspaceDir(t, w)
    49  	f, err := filer.NewWorkspaceFilesClient(w, tmpdir)
    50  	require.NoError(t, err)
    51  
    52  	// Write file to workspace
    53  	contents := "#!/usr/bin/bash\necho hello, world\n"
    54  	err = f.Write(ctx, "file-a", strings.NewReader(contents))
    55  	require.NoError(t, err)
    56  
    57  	// Run export
    58  	stdout, stderr := RequireSuccessfulRun(t, "workspace", "export", path.Join(tmpdir, "file-a"))
    59  	assert.Equal(t, contents, stdout.String())
    60  	assert.Equal(t, "", stderr.String())
    61  }
    62  
    63  func setupWorkspaceImportExportTest(t *testing.T) (context.Context, filer.Filer, string) {
    64  	t.Log(GetEnvOrSkipTest(t, "CLOUD_ENV"))
    65  
    66  	ctx := context.Background()
    67  	w := databricks.Must(databricks.NewWorkspaceClient())
    68  	tmpdir := temporaryWorkspaceDir(t, w)
    69  	f, err := filer.NewWorkspaceFilesClient(w, tmpdir)
    70  	require.NoError(t, err)
    71  
    72  	// Check if we can use this API here, skip test if we cannot.
    73  	_, err = f.Read(ctx, "we_use_this_call_to_test_if_this_api_is_enabled")
    74  	var aerr *apierr.APIError
    75  	if errors.As(err, &aerr) && aerr.StatusCode == http.StatusBadRequest {
    76  		t.Skip(aerr.Message)
    77  	}
    78  
    79  	return ctx, f, tmpdir
    80  }
    81  
    82  // TODO: add tests for the progress event output logs: https://github.com/databricks/cli/issues/447
    83  func assertLocalFileContents(t *testing.T, path string, content string) {
    84  	require.FileExists(t, path)
    85  	b, err := os.ReadFile(path)
    86  	require.NoError(t, err)
    87  	assert.Contains(t, string(b), content)
    88  }
    89  
    90  func assertFilerFileContents(t *testing.T, ctx context.Context, f filer.Filer, path string, content string) {
    91  	r, err := f.Read(ctx, path)
    92  	require.NoError(t, err)
    93  	b, err := io.ReadAll(r)
    94  	require.NoError(t, err)
    95  	assert.Contains(t, string(b), content)
    96  }
    97  
    98  func TestAccExportDir(t *testing.T) {
    99  	ctx, f, sourceDir := setupWorkspaceImportExportTest(t)
   100  	targetDir := t.TempDir()
   101  
   102  	var err error
   103  
   104  	// Write test data to the workspace
   105  	err = f.Write(ctx, "file-a", strings.NewReader("abc"))
   106  	require.NoError(t, err)
   107  	err = f.Write(ctx, "pyNotebook.py", strings.NewReader("# Databricks notebook source"))
   108  	require.NoError(t, err)
   109  	err = f.Write(ctx, "sqlNotebook.sql", strings.NewReader("-- Databricks notebook source"))
   110  	require.NoError(t, err)
   111  	err = f.Write(ctx, "scalaNotebook.scala", strings.NewReader("// Databricks notebook source"))
   112  	require.NoError(t, err)
   113  	err = f.Write(ctx, "rNotebook.r", strings.NewReader("# Databricks notebook source"))
   114  	require.NoError(t, err)
   115  	err = f.Write(ctx, "a/b/c/file-b", strings.NewReader("def"), filer.CreateParentDirectories)
   116  	require.NoError(t, err)
   117  
   118  	// Run Export
   119  	RequireSuccessfulRun(t, "workspace", "export-dir", sourceDir, targetDir)
   120  
   121  	// Assert files were exported
   122  	assertLocalFileContents(t, filepath.Join(targetDir, "file-a"), "abc")
   123  	assertLocalFileContents(t, filepath.Join(targetDir, "pyNotebook.py"), "# Databricks notebook source")
   124  	assertLocalFileContents(t, filepath.Join(targetDir, "sqlNotebook.sql"), "-- Databricks notebook source")
   125  	assertLocalFileContents(t, filepath.Join(targetDir, "rNotebook.r"), "# Databricks notebook source")
   126  	assertLocalFileContents(t, filepath.Join(targetDir, "scalaNotebook.scala"), "// Databricks notebook source")
   127  	assertLocalFileContents(t, filepath.Join(targetDir, "a/b/c/file-b"), "def")
   128  }
   129  
   130  func TestAccExportDirDoesNotOverwrite(t *testing.T) {
   131  	ctx, f, sourceDir := setupWorkspaceImportExportTest(t)
   132  	targetDir := t.TempDir()
   133  
   134  	var err error
   135  
   136  	// Write remote file
   137  	err = f.Write(ctx, "file-a", strings.NewReader("content from workspace"))
   138  	require.NoError(t, err)
   139  
   140  	// Write local file
   141  	err = os.WriteFile(filepath.Join(targetDir, "file-a"), []byte("local content"), os.ModePerm)
   142  	require.NoError(t, err)
   143  
   144  	// Run Export
   145  	RequireSuccessfulRun(t, "workspace", "export-dir", sourceDir, targetDir)
   146  
   147  	// Assert file is not overwritten
   148  	assertLocalFileContents(t, filepath.Join(targetDir, "file-a"), "local content")
   149  }
   150  
   151  func TestAccExportDirWithOverwriteFlag(t *testing.T) {
   152  	ctx, f, sourceDir := setupWorkspaceImportExportTest(t)
   153  	targetDir := t.TempDir()
   154  
   155  	var err error
   156  
   157  	// Write remote file
   158  	err = f.Write(ctx, "file-a", strings.NewReader("content from workspace"))
   159  	require.NoError(t, err)
   160  
   161  	// Write local file
   162  	err = os.WriteFile(filepath.Join(targetDir, "file-a"), []byte("local content"), os.ModePerm)
   163  	require.NoError(t, err)
   164  
   165  	// Run Export
   166  	RequireSuccessfulRun(t, "workspace", "export-dir", sourceDir, targetDir, "--overwrite")
   167  
   168  	// Assert file has been overwritten
   169  	assertLocalFileContents(t, filepath.Join(targetDir, "file-a"), "content from workspace")
   170  }
   171  
   172  // TODO: Add assertions on progress logs for workspace import-dir command. https://github.com/databricks/cli/issues/455
   173  func TestAccImportDir(t *testing.T) {
   174  	ctx, workspaceFiler, targetDir := setupWorkspaceImportExportTest(t)
   175  	RequireSuccessfulRun(t, "workspace", "import-dir", "./testdata/import_dir", targetDir, "--log-level=debug")
   176  
   177  	// Assert files are imported
   178  	assertFilerFileContents(t, ctx, workspaceFiler, "file-a", "hello, world")
   179  	assertFilerFileContents(t, ctx, workspaceFiler, "a/b/c/file-b", "file-in-dir")
   180  	assertFilerFileContents(t, ctx, workspaceFiler, "pyNotebook", "# Databricks notebook source\nprint(\"python\")")
   181  	assertFilerFileContents(t, ctx, workspaceFiler, "sqlNotebook", "-- Databricks notebook source\nSELECT \"sql\"")
   182  	assertFilerFileContents(t, ctx, workspaceFiler, "rNotebook", "# Databricks notebook source\nprint(\"r\")")
   183  	assertFilerFileContents(t, ctx, workspaceFiler, "scalaNotebook", "// Databricks notebook source\nprintln(\"scala\")")
   184  	assertFilerFileContents(t, ctx, workspaceFiler, "jupyterNotebook", "# Databricks notebook source\nprint(\"jupyter\")")
   185  }
   186  
   187  func TestAccImportDirDoesNotOverwrite(t *testing.T) {
   188  	ctx, workspaceFiler, targetDir := setupWorkspaceImportExportTest(t)
   189  	var err error
   190  
   191  	// create preexisting files in the workspace
   192  	err = workspaceFiler.Write(ctx, "file-a", strings.NewReader("old file"))
   193  	require.NoError(t, err)
   194  	err = workspaceFiler.Write(ctx, "pyNotebook.py", strings.NewReader("# Databricks notebook source\nprint(\"old notebook\")"))
   195  	require.NoError(t, err)
   196  
   197  	// Assert contents of pre existing files
   198  	assertFilerFileContents(t, ctx, workspaceFiler, "file-a", "old file")
   199  	assertFilerFileContents(t, ctx, workspaceFiler, "pyNotebook", "# Databricks notebook source\nprint(\"old notebook\")")
   200  
   201  	RequireSuccessfulRun(t, "workspace", "import-dir", "./testdata/import_dir", targetDir)
   202  
   203  	// Assert files are imported
   204  	assertFilerFileContents(t, ctx, workspaceFiler, "a/b/c/file-b", "file-in-dir")
   205  	assertFilerFileContents(t, ctx, workspaceFiler, "sqlNotebook", "-- Databricks notebook source\nSELECT \"sql\"")
   206  	assertFilerFileContents(t, ctx, workspaceFiler, "rNotebook", "# Databricks notebook source\nprint(\"r\")")
   207  	assertFilerFileContents(t, ctx, workspaceFiler, "scalaNotebook", "// Databricks notebook source\nprintln(\"scala\")")
   208  	assertFilerFileContents(t, ctx, workspaceFiler, "jupyterNotebook", "# Databricks notebook source\nprint(\"jupyter\")")
   209  
   210  	// Assert pre existing files are not changed
   211  	assertFilerFileContents(t, ctx, workspaceFiler, "file-a", "old file")
   212  	assertFilerFileContents(t, ctx, workspaceFiler, "pyNotebook", "# Databricks notebook source\nprint(\"old notebook\")")
   213  }
   214  
   215  func TestAccImportDirWithOverwriteFlag(t *testing.T) {
   216  	ctx, workspaceFiler, targetDir := setupWorkspaceImportExportTest(t)
   217  	var err error
   218  
   219  	// create preexisting files in the workspace
   220  	err = workspaceFiler.Write(ctx, "file-a", strings.NewReader("old file"))
   221  	require.NoError(t, err)
   222  	err = workspaceFiler.Write(ctx, "pyNotebook.py", strings.NewReader("# Databricks notebook source\nprint(\"old notebook\")"))
   223  	require.NoError(t, err)
   224  
   225  	// Assert contents of pre existing files
   226  	assertFilerFileContents(t, ctx, workspaceFiler, "file-a", "old file")
   227  	assertFilerFileContents(t, ctx, workspaceFiler, "pyNotebook", "# Databricks notebook source\nprint(\"old notebook\")")
   228  
   229  	RequireSuccessfulRun(t, "workspace", "import-dir", "./testdata/import_dir", targetDir, "--overwrite")
   230  
   231  	// Assert files are imported
   232  	assertFilerFileContents(t, ctx, workspaceFiler, "a/b/c/file-b", "file-in-dir")
   233  	assertFilerFileContents(t, ctx, workspaceFiler, "sqlNotebook", "-- Databricks notebook source\nSELECT \"sql\"")
   234  	assertFilerFileContents(t, ctx, workspaceFiler, "rNotebook", "# Databricks notebook source\nprint(\"r\")")
   235  	assertFilerFileContents(t, ctx, workspaceFiler, "scalaNotebook", "// Databricks notebook source\nprintln(\"scala\")")
   236  	assertFilerFileContents(t, ctx, workspaceFiler, "jupyterNotebook", "# Databricks notebook source\nprint(\"jupyter\")")
   237  
   238  	// Assert pre existing files are overwritten
   239  	assertFilerFileContents(t, ctx, workspaceFiler, "file-a", "hello, world")
   240  	assertFilerFileContents(t, ctx, workspaceFiler, "pyNotebook", "# Databricks notebook source\nprint(\"python\")")
   241  }