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 }