databricks-cli/cmd/auth/describe_test.go

package auth

import (
	"context"
	"fmt"
	"testing"

	"github.com/databricks/cli/cmd/root"
	"github.com/databricks/databricks-sdk-go/config"
	"github.com/databricks/databricks-sdk-go/experimental/mocks"
	"github.com/databricks/databricks-sdk-go/service/iam"
	"github.com/spf13/cobra"
	"github.com/stretchr/testify/mock"
	"github.com/stretchr/testify/require"
)

func TestGetWorkspaceAuthStatus(t *testing.T) {
	ctx := context.Background()
	m := mocks.NewMockWorkspaceClient(t)
	ctx = root.SetWorkspaceClient(ctx, m.WorkspaceClient)

	cmd := &cobra.Command{}
	cmd.SetContext(ctx)

	showSensitive := false

	currentUserApi := m.GetMockCurrentUserAPI()
	currentUserApi.EXPECT().Me(mock.Anything).Return(&iam.User{
		UserName: "test-user",
	}, nil)

	cmd.Flags().String("host", "", "")
	cmd.Flags().String("profile", "", "")
	err := cmd.Flag("profile").Value.Set("my-profile")
	require.NoError(t, err)
	cmd.Flag("profile").Changed = true

	cfg := &config.Config{
		Profile: "my-profile",
	}
	m.WorkspaceClient.Config = cfg
	t.Setenv("DATABRICKS_AUTH_TYPE", "azure-cli")
	err = config.ConfigAttributes.Configure(cfg)
	require.NoError(t, err)

	status, err := getAuthStatus(cmd, []string{}, showSensitive, func(cmd *cobra.Command, args []string) (*config.Config, bool, error) {
		err := config.ConfigAttributes.ResolveFromStringMap(cfg, map[string]string{
			"host":      "https://test.com",
			"token":     "test-token",
			"auth_type": "azure-cli",
		})
		require.NoError(t, err)
		return cfg, false, nil
	})
	require.NoError(t, err)
	require.NotNil(t, status)
	require.Equal(t, "success", status.Status)
	require.Equal(t, "test-user", status.Username)
	require.Equal(t, "https://test.com", status.Details.Host)
	require.Equal(t, "azure-cli", status.Details.AuthType)

	require.Equal(t, "azure-cli", status.Details.Configuration["auth_type"].Value)
	require.Equal(t, "DATABRICKS_AUTH_TYPE environment variable", status.Details.Configuration["auth_type"].Source.String())
	require.False(t, status.Details.Configuration["auth_type"].AuthTypeMismatch)

	require.Equal(t, "********", status.Details.Configuration["token"].Value)
	require.Equal(t, "dynamic configuration", status.Details.Configuration["token"].Source.String())
	require.True(t, status.Details.Configuration["token"].AuthTypeMismatch)

	require.Equal(t, "my-profile", status.Details.Configuration["profile"].Value)
	require.Equal(t, "--profile flag", status.Details.Configuration["profile"].Source.String())
	require.False(t, status.Details.Configuration["profile"].AuthTypeMismatch)
}

func TestGetWorkspaceAuthStatusError(t *testing.T) {
	ctx := context.Background()
	m := mocks.NewMockWorkspaceClient(t)
	ctx = root.SetWorkspaceClient(ctx, m.WorkspaceClient)

	cmd := &cobra.Command{}
	cmd.SetContext(ctx)

	showSensitive := false

	cmd.Flags().String("host", "", "")
	cmd.Flags().String("profile", "", "")
	err := cmd.Flag("profile").Value.Set("my-profile")
	require.NoError(t, err)
	cmd.Flag("profile").Changed = true

	cfg := &config.Config{
		Profile: "my-profile",
	}
	m.WorkspaceClient.Config = cfg
	t.Setenv("DATABRICKS_AUTH_TYPE", "azure-cli")
	err = config.ConfigAttributes.Configure(cfg)
	require.NoError(t, err)

	status, err := getAuthStatus(cmd, []string{}, showSensitive, func(cmd *cobra.Command, args []string) (*config.Config, bool, error) {
		err = config.ConfigAttributes.ResolveFromStringMap(cfg, map[string]string{
			"host":      "https://test.com",
			"token":     "test-token",
			"auth_type": "azure-cli",
		})
		return cfg, false, fmt.Errorf("auth error")
	})
	require.NoError(t, err)
	require.NotNil(t, status)
	require.Equal(t, "error", status.Status)

	require.Equal(t, "azure-cli", status.Details.Configuration["auth_type"].Value)
	require.Equal(t, "DATABRICKS_AUTH_TYPE environment variable", status.Details.Configuration["auth_type"].Source.String())
	require.False(t, status.Details.Configuration["auth_type"].AuthTypeMismatch)

	require.Equal(t, "********", status.Details.Configuration["token"].Value)
	require.Equal(t, "dynamic configuration", status.Details.Configuration["token"].Source.String())
	require.True(t, status.Details.Configuration["token"].AuthTypeMismatch)

	require.Equal(t, "my-profile", status.Details.Configuration["profile"].Value)
	require.Equal(t, "--profile flag", status.Details.Configuration["profile"].Source.String())
	require.False(t, status.Details.Configuration["profile"].AuthTypeMismatch)
}

func TestGetWorkspaceAuthStatusSensitive(t *testing.T) {
	ctx := context.Background()
	m := mocks.NewMockWorkspaceClient(t)
	ctx = root.SetWorkspaceClient(ctx, m.WorkspaceClient)

	cmd := &cobra.Command{}
	cmd.SetContext(ctx)

	showSensitive := true

	cmd.Flags().String("host", "", "")
	cmd.Flags().String("profile", "", "")
	err := cmd.Flag("profile").Value.Set("my-profile")
	require.NoError(t, err)
	cmd.Flag("profile").Changed = true

	cfg := &config.Config{
		Profile: "my-profile",
	}
	m.WorkspaceClient.Config = cfg
	t.Setenv("DATABRICKS_AUTH_TYPE", "azure-cli")
	err = config.ConfigAttributes.Configure(cfg)
	require.NoError(t, err)

	status, err := getAuthStatus(cmd, []string{}, showSensitive, func(cmd *cobra.Command, args []string) (*config.Config, bool, error) {
		err = config.ConfigAttributes.ResolveFromStringMap(cfg, map[string]string{
			"host":      "https://test.com",
			"token":     "test-token",
			"auth_type": "azure-cli",
		})
		return cfg, false, fmt.Errorf("auth error")
	})
	require.NoError(t, err)
	require.NotNil(t, status)
	require.Equal(t, "error", status.Status)

	require.Equal(t, "azure-cli", status.Details.Configuration["auth_type"].Value)
	require.Equal(t, "DATABRICKS_AUTH_TYPE environment variable", status.Details.Configuration["auth_type"].Source.String())
	require.False(t, status.Details.Configuration["auth_type"].AuthTypeMismatch)

	require.Equal(t, "test-token", status.Details.Configuration["token"].Value)
	require.Equal(t, "dynamic configuration", status.Details.Configuration["token"].Source.String())
	require.True(t, status.Details.Configuration["token"].AuthTypeMismatch)
}

func TestGetAccountAuthStatus(t *testing.T) {
	ctx := context.Background()
	m := mocks.NewMockAccountClient(t)
	ctx = root.SetAccountClient(ctx, m.AccountClient)

	cmd := &cobra.Command{}
	cmd.SetContext(ctx)

	showSensitive := false

	cmd.Flags().String("host", "", "")
	cmd.Flags().String("profile", "", "")
	err := cmd.Flag("profile").Value.Set("my-profile")
	require.NoError(t, err)
	cmd.Flag("profile").Changed = true

	cfg := &config.Config{
		Profile: "my-profile",
	}
	m.AccountClient.Config = cfg
	t.Setenv("DATABRICKS_AUTH_TYPE", "azure-cli")
	err = config.ConfigAttributes.Configure(cfg)
	require.NoError(t, err)

	wsApi := m.GetMockWorkspacesAPI()
	wsApi.EXPECT().List(mock.Anything).Return(nil, nil)

	status, err := getAuthStatus(cmd, []string{}, showSensitive, func(cmd *cobra.Command, args []string) (*config.Config, bool, error) {
		err = config.ConfigAttributes.ResolveFromStringMap(cfg, map[string]string{
			"account_id": "test-account-id",
			"username":   "test-user",
			"host":       "https://test.com",
			"token":      "test-token",
			"auth_type":  "azure-cli",
		})
		return cfg, true, nil
	})
	require.NoError(t, err)
	require.NotNil(t, status)
	require.Equal(t, "success", status.Status)

	require.Equal(t, "test-user", status.Username)
	require.Equal(t, "https://test.com", status.Details.Host)
	require.Equal(t, "azure-cli", status.Details.AuthType)
	require.Equal(t, "test-account-id", status.AccountID)

	require.Equal(t, "azure-cli", status.Details.Configuration["auth_type"].Value)
	require.Equal(t, "DATABRICKS_AUTH_TYPE environment variable", status.Details.Configuration["auth_type"].Source.String())
	require.False(t, status.Details.Configuration["auth_type"].AuthTypeMismatch)

	require.Equal(t, "********", status.Details.Configuration["token"].Value)
	require.Equal(t, "dynamic configuration", status.Details.Configuration["token"].Source.String())
	require.True(t, status.Details.Configuration["token"].AuthTypeMismatch)

	require.Equal(t, "my-profile", status.Details.Configuration["profile"].Value)
	require.Equal(t, "--profile flag", status.Details.Configuration["profile"].Source.String())
	require.False(t, status.Details.Configuration["profile"].AuthTypeMismatch)
}