mirror of https://github.com/databricks/cli.git
318 lines
8.2 KiB
Go
318 lines
8.2 KiB
Go
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",
|
|
|
|
// Force SDK to not try and lookup the tenant ID from the host.
|
|
// The host above is invalid and will not be reachable.
|
|
AzureTenantID: "nonempty",
|
|
})
|
|
})
|
|
|
|
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/",
|
|
|
|
// Force SDK to not try and lookup the tenant ID from the host.
|
|
// The host above is invalid and will not be reachable.
|
|
AzureTenantID: "nonempty",
|
|
})
|
|
})
|
|
|
|
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",
|
|
})
|
|
})
|
|
}
|
|
|
|
func TestMustAccountClientWorksWithDatabricksCfg(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
|
|
`),
|
|
0755)
|
|
require.NoError(t, err)
|
|
|
|
cmd := New(context.Background())
|
|
|
|
t.Setenv("DATABRICKS_CONFIG_FILE", configFile)
|
|
err = MustAccountClient(cmd, []string{})
|
|
require.NoError(t, err)
|
|
}
|
|
|
|
func TestMustAccountClientWorksWithNoDatabricksCfgButEnvironmentVariables(t *testing.T) {
|
|
testutil.CleanupEnvironment(t)
|
|
|
|
ctx, tt := cmdio.SetupTest(context.Background())
|
|
t.Cleanup(tt.Done)
|
|
cmd := New(ctx)
|
|
t.Setenv("DATABRICKS_HOST", "https://accounts.azuredatabricks.net/")
|
|
t.Setenv("DATABRICKS_TOKEN", "foobar")
|
|
t.Setenv("DATABRICKS_ACCOUNT_ID", "1111")
|
|
|
|
err := MustAccountClient(cmd, []string{})
|
|
require.NoError(t, err)
|
|
}
|
|
|
|
func TestMustAccountClientErrorsWithNoDatabricksCfg(t *testing.T) {
|
|
testutil.CleanupEnvironment(t)
|
|
|
|
ctx, tt := cmdio.SetupTest(context.Background())
|
|
t.Cleanup(tt.Done)
|
|
cmd := New(ctx)
|
|
|
|
err := MustAccountClient(cmd, []string{})
|
|
require.ErrorContains(t, err, "no configuration file found at")
|
|
}
|
|
|
|
func TestMustAnyClientCanCreateWorkspaceClient(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
|
|
`),
|
|
0755)
|
|
require.NoError(t, err)
|
|
|
|
ctx, tt := cmdio.SetupTest(context.Background())
|
|
t.Cleanup(tt.Done)
|
|
cmd := New(ctx)
|
|
|
|
t.Setenv("DATABRICKS_CONFIG_FILE", configFile)
|
|
isAccount, err := MustAnyClient(cmd, []string{})
|
|
require.False(t, isAccount)
|
|
require.NoError(t, err)
|
|
|
|
w := WorkspaceClient(cmd.Context())
|
|
require.NotNil(t, w)
|
|
}
|
|
|
|
func TestMustAnyClientCanCreateAccountClient(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
|
|
`),
|
|
0755)
|
|
require.NoError(t, err)
|
|
|
|
ctx, tt := cmdio.SetupTest(context.Background())
|
|
t.Cleanup(tt.Done)
|
|
cmd := New(ctx)
|
|
|
|
t.Setenv("DATABRICKS_CONFIG_FILE", configFile)
|
|
isAccount, err := MustAnyClient(cmd, []string{})
|
|
require.NoError(t, err)
|
|
require.True(t, isAccount)
|
|
|
|
a := AccountClient(cmd.Context())
|
|
require.NotNil(t, a)
|
|
}
|
|
|
|
func TestMustAnyClientWithEmptyDatabricksCfg(t *testing.T) {
|
|
testutil.CleanupEnvironment(t)
|
|
|
|
dir := t.TempDir()
|
|
configFile := filepath.Join(dir, ".databrickscfg")
|
|
err := os.WriteFile(
|
|
configFile,
|
|
[]byte(""), // empty file
|
|
0755)
|
|
require.NoError(t, err)
|
|
|
|
ctx, tt := cmdio.SetupTest(context.Background())
|
|
t.Cleanup(tt.Done)
|
|
cmd := New(ctx)
|
|
|
|
t.Setenv("DATABRICKS_CONFIG_FILE", configFile)
|
|
|
|
_, err = MustAnyClient(cmd, []string{})
|
|
require.ErrorContains(t, err, "does not contain account profiles")
|
|
}
|