package root import ( "context" "os" "path/filepath" "testing" "time" "github.com/databricks/cli/internal/testutil" "github.com/databricks/cli/libs/cmdio" "github.com/databricks/databricks-sdk-go/config" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) func TestEmptyHttpRequest(t *testing.T) { ctx, _ := context.WithCancel(context.Background()) req := emptyHttpRequest(ctx) assert.Equal(t, req.Context(), ctx) } type promptFn func(ctx context.Context, cfg *config.Config, retry bool) (any, error) var accountPromptFn = func(ctx context.Context, cfg *config.Config, retry bool) (any, error) { return accountClientOrPrompt(ctx, cfg, retry) } var workspacePromptFn = func(ctx context.Context, cfg *config.Config, retry bool) (any, error) { return workspaceClientOrPrompt(ctx, cfg, retry) } func expectPrompts(t *testing.T, fn promptFn, config *config.Config) { ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second) defer cancel() // Channel to pass errors from the prompting function back to the test. errch := make(chan error, 1) ctx, io := cmdio.SetupTest(ctx) go func() { defer close(errch) defer cancel() _, err := fn(ctx, config, true) errch <- err }() // Expect a prompt line, _, err := io.Stderr.ReadLine() if assert.NoError(t, err, "Expected to read a line from stderr") { assert.Contains(t, string(line), "Search:") } else { // If there was an error reading from stderr, the prompting function must have terminated early. assert.NoError(t, <-errch) } } func expectReturns(t *testing.T, fn promptFn, config *config.Config) { ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second) defer cancel() ctx, _ = cmdio.SetupTest(ctx) client, err := fn(ctx, config, true) require.NoError(t, err) require.NotNil(t, client) } func TestAccountClientOrPrompt(t *testing.T) { testutil.CleanupEnvironment(t) dir := t.TempDir() configFile := filepath.Join(dir, ".databrickscfg") err := os.WriteFile( configFile, []byte(` [account-1111] host = https://accounts.azuredatabricks.net/ account_id = 1111 token = foobar [account-1112] host = https://accounts.azuredatabricks.net/ account_id = 1112 token = foobar `), 0755) require.NoError(t, err) t.Setenv("DATABRICKS_CONFIG_FILE", configFile) t.Setenv("PATH", "/nothing") t.Run("Prompt if nothing is specified", func(t *testing.T) { expectPrompts(t, accountPromptFn, &config.Config{}) }) t.Run("Prompt if a workspace host is specified", func(t *testing.T) { expectPrompts(t, accountPromptFn, &config.Config{ Host: "https://adb-1234567.89.azuredatabricks.net/", AccountID: "1234", Token: "foobar", }) }) t.Run("Prompt if account ID is not specified", func(t *testing.T) { expectPrompts(t, accountPromptFn, &config.Config{ Host: "https://accounts.azuredatabricks.net/", Token: "foobar", }) }) t.Run("Prompt if no credential provider can be configured", func(t *testing.T) { expectPrompts(t, accountPromptFn, &config.Config{ Host: "https://accounts.azuredatabricks.net/", AccountID: "1234", }) }) t.Run("Returns if configuration is valid", func(t *testing.T) { expectReturns(t, accountPromptFn, &config.Config{ Host: "https://accounts.azuredatabricks.net/", AccountID: "1234", Token: "foobar", }) }) t.Run("Returns if a valid profile is specified", func(t *testing.T) { expectReturns(t, accountPromptFn, &config.Config{ Profile: "account-1111", }) }) } func TestWorkspaceClientOrPrompt(t *testing.T) { testutil.CleanupEnvironment(t) dir := t.TempDir() configFile := filepath.Join(dir, ".databrickscfg") err := os.WriteFile( configFile, []byte(` [workspace-1111] host = https://adb-1111.11.azuredatabricks.net/ token = foobar [workspace-1112] host = https://adb-1112.12.azuredatabricks.net/ token = foobar `), 0755) require.NoError(t, err) t.Setenv("DATABRICKS_CONFIG_FILE", configFile) t.Setenv("PATH", "/nothing") t.Run("Prompt if nothing is specified", func(t *testing.T) { expectPrompts(t, workspacePromptFn, &config.Config{}) }) t.Run("Prompt if an account host is specified", func(t *testing.T) { expectPrompts(t, workspacePromptFn, &config.Config{ Host: "https://accounts.azuredatabricks.net/", AccountID: "1234", Token: "foobar", }) }) t.Run("Prompt if no credential provider can be configured", func(t *testing.T) { expectPrompts(t, workspacePromptFn, &config.Config{ Host: "https://adb-1111.11.azuredatabricks.net/", }) }) t.Run("Returns if configuration is valid", func(t *testing.T) { expectReturns(t, workspacePromptFn, &config.Config{ Host: "https://adb-1111.11.azuredatabricks.net/", Token: "foobar", }) }) t.Run("Returns if a valid profile is specified", func(t *testing.T) { expectReturns(t, workspacePromptFn, &config.Config{ Profile: "workspace-1111", }) }) }