package internal import ( "bufio" "bytes" "context" "fmt" "math/rand" "os" "path/filepath" "strings" "testing" "time" "github.com/databricks/cli/cmd/root" _ "github.com/databricks/cli/cmd/version" "github.com/databricks/databricks-sdk-go" "github.com/databricks/databricks-sdk-go/apierr" "github.com/databricks/databricks-sdk-go/service/workspace" "github.com/stretchr/testify/require" _ "github.com/databricks/cli/cmd/workspace" ) const charset = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789" // GetEnvOrSkipTest proceeds with test only with that env variable func GetEnvOrSkipTest(t *testing.T, name string) string { value := os.Getenv(name) if value == "" { t.Skipf("Environment variable %s is missing", name) } return value } // RandomName gives random name with optional prefix. e.g. qa.RandomName("tf-") func RandomName(prefix ...string) string { rand.Seed(time.Now().UnixNano()) randLen := 12 b := make([]byte, randLen) for i := range b { b[i] = charset[rand.Intn(randLen)] } if len(prefix) > 0 { return fmt.Sprintf("%s%s", strings.Join(prefix, ""), b) } return string(b) } // Helper for running the root command in the background. // It ensures that the background goroutine terminates upon // test completion through cancelling the command context. type cobraTestRunner struct { *testing.T args []string stdout bytes.Buffer stderr bytes.Buffer errch <-chan error } func (t *cobraTestRunner) RunBackground() { root := root.RootCmd root.SetOut(&t.stdout) root.SetErr(&t.stderr) root.SetArgs(t.args) errch := make(chan error) ctx, cancel := context.WithCancel(context.Background()) // Run command in background. go func() { cmd, err := root.ExecuteContextC(ctx) if err != nil { t.Logf("Error running command: %s", err) } if t.stdout.Len() > 0 { // Make a copy of the buffer such that it remains "unread". scanner := bufio.NewScanner(bytes.NewBuffer(t.stdout.Bytes())) for scanner.Scan() { t.Logf("[databricks stdout]: %s", scanner.Text()) } } if t.stderr.Len() > 0 { // Make a copy of the buffer such that it remains "unread". scanner := bufio.NewScanner(bytes.NewBuffer(t.stderr.Bytes())) for scanner.Scan() { t.Logf("[databricks stderr]: %s", scanner.Text()) } } // Reset context on command for the next test. // These commands are globals so we have to clean up to the best of our ability after each run. // See https://github.com/spf13/cobra/blob/a6f198b635c4b18fff81930c40d464904e55b161/command.go#L1062-L1066 //lint:ignore SA1012 cobra sets the context and doesn't clear it cmd.SetContext(nil) // Make caller aware of error. errch <- err close(errch) }() // Ensure command terminates upon test completion (success or failure). t.Cleanup(func() { // Signal termination of command. cancel() // Wait for goroutine to finish. <-errch }) t.errch = errch } func (t *cobraTestRunner) Run() (bytes.Buffer, bytes.Buffer, error) { t.RunBackground() err := <-t.errch return t.stdout, t.stderr, err } // Like [require.Eventually] but errors if the underlying command has failed. func (c *cobraTestRunner) Eventually(condition func() bool, waitFor time.Duration, tick time.Duration, msgAndArgs ...interface{}) { ch := make(chan bool, 1) timer := time.NewTimer(waitFor) defer timer.Stop() ticker := time.NewTicker(tick) defer ticker.Stop() for tick := ticker.C; ; { select { case err := <-c.errch: require.Fail(c, "Command failed", err) return case <-timer.C: require.Fail(c, "Condition never satisfied", msgAndArgs...) return case <-tick: tick = nil go func() { ch <- condition() }() case v := <-ch: if v { return } tick = ticker.C } } } func NewCobraTestRunner(t *testing.T, args ...string) *cobraTestRunner { return &cobraTestRunner{ T: t, args: args, } } func RequireSuccessfulRun(t *testing.T, args ...string) (bytes.Buffer, bytes.Buffer) { c := NewCobraTestRunner(t, args...) stdout, stderr, err := c.Run() require.NoError(t, err) return stdout, stderr } func RequireErrorRun(t *testing.T, args ...string) (bytes.Buffer, bytes.Buffer, error) { c := NewCobraTestRunner(t, args...) stdout, stderr, err := c.Run() require.Error(t, err) return stdout, stderr, err } func writeFile(t *testing.T, name string, body string) string { f, err := os.Create(filepath.Join(t.TempDir(), name)) require.NoError(t, err) _, err = f.WriteString(body) require.NoError(t, err) f.Close() return f.Name() } func TemporaryWorkspaceDir(t *testing.T, w *databricks.WorkspaceClient) string { ctx := context.Background() me, err := w.CurrentUser.Me(ctx) require.NoError(t, err) path := fmt.Sprintf("/Users/%s/%s", me.UserName, RandomName("integration-test-wsfs-")) // Ensure directory exists, but doesn't exist YET! // Otherwise we could inadvertently remove a directory that already exists on cleanup. t.Logf("mkdir %s", path) err = w.Workspace.MkdirsByPath(ctx, path) require.NoError(t, err) // Remove test directory on test completion. t.Cleanup(func() { t.Logf("rm -rf %s", path) err := w.Workspace.Delete(ctx, workspace.Delete{ Path: path, Recursive: true, }) if err == nil || apierr.IsMissing(err) { return } t.Logf("unable to remove temporary workspace directory %s: %#v", path, err) }) return path }