databricks-cli/libs/auth/env_test.go

package auth

import (
	"testing"

	"github.com/databricks/cli/internal/testutil"
	"github.com/databricks/databricks-sdk-go/config"
	"github.com/stretchr/testify/assert"
)

func TestAuthEnv(t *testing.T) {
	in := &config.Config{
		Profile:            "thisshouldbeignored",
		Host:               "https://test.com",
		Token:              "test-token",
		Password:           "test-password",
		MetadataServiceURL: "http://somurl.com",

		AzureUseMSI:       true,
		AzureTenantID:     "test-tenant-id",
		AzureClientID:     "test-client-id",
		AzureClientSecret: "test-client-secret",

		ActionsIDTokenRequestToken: "test-actions-id-token-request-token",
	}

	expected := map[string]string{
		"DATABRICKS_HOST":                 "https://test.com",
		"DATABRICKS_TOKEN":                "test-token",
		"DATABRICKS_PASSWORD":             "test-password",
		"DATABRICKS_METADATA_SERVICE_URL": "http://somurl.com",

		"ARM_USE_MSI":       "true",
		"ARM_TENANT_ID":     "test-tenant-id",
		"ARM_CLIENT_ID":     "test-client-id",
		"ARM_CLIENT_SECRET": "test-client-secret",

		"ACTIONS_ID_TOKEN_REQUEST_TOKEN": "test-actions-id-token-request-token",
	}

	out := Env(in)
	assert.Equal(t, expected, out)
}

func TestGetEnvFor(t *testing.T) {
	tcases := []struct {
		name     string
		expected string
	}{
		// Generic attributes.
		{"host", "DATABRICKS_HOST"},
		{"profile", "DATABRICKS_CONFIG_PROFILE"},
		{"auth_type", "DATABRICKS_AUTH_TYPE"},
		{"metadata_service_url", "DATABRICKS_METADATA_SERVICE_URL"},

		// OAuth specific attributes.
		{"client_id", "DATABRICKS_CLIENT_ID"},

		// Google specific attributes.
		{"google_service_account", "DATABRICKS_GOOGLE_SERVICE_ACCOUNT"},

		// Azure specific attributes.
		{"azure_workspace_resource_id", "DATABRICKS_AZURE_RESOURCE_ID"},
		{"azure_use_msi", "ARM_USE_MSI"},
		{"azure_client_id", "ARM_CLIENT_ID"},
		{"azure_tenant_id", "ARM_TENANT_ID"},
		{"azure_environment", "ARM_ENVIRONMENT"},
		{"azure_login_app_id", "DATABRICKS_AZURE_LOGIN_APP_ID"},
	}

	for _, tcase := range tcases {
		t.Run(tcase.name, func(t *testing.T) {
			out, ok := GetEnvFor(tcase.name)
			assert.True(t, ok)
			assert.Equal(t, tcase.expected, out)
		})
	}

	out, ok := GetEnvFor("notfound")
	assert.False(t, ok)
	assert.Empty(t, out)
}

func TestAuthEnvVars(t *testing.T) {
	expected := []string{
		"DATABRICKS_HOST",
		"DATABRICKS_CLUSTER_ID",
		"DATABRICKS_WAREHOUSE_ID",
		"DATABRICKS_SERVERLESS_COMPUTE_ID",
		"DATABRICKS_METADATA_SERVICE_URL",
		"DATABRICKS_ACCOUNT_ID",
		"DATABRICKS_TOKEN",
		"DATABRICKS_USERNAME",
		"DATABRICKS_PASSWORD",
		"DATABRICKS_CONFIG_PROFILE",
		"DATABRICKS_CONFIG_FILE",
		"DATABRICKS_GOOGLE_SERVICE_ACCOUNT",
		"GOOGLE_CREDENTIALS",
		"DATABRICKS_AZURE_RESOURCE_ID",
		"ARM_USE_MSI",
		"ARM_CLIENT_SECRET",
		"ARM_CLIENT_ID",
		"ARM_TENANT_ID",
		"ACTIONS_ID_TOKEN_REQUEST_URL",
		"ACTIONS_ID_TOKEN_REQUEST_TOKEN",
		"ARM_ENVIRONMENT",
		"DATABRICKS_AZURE_LOGIN_APP_ID",
		"DATABRICKS_CLIENT_ID",
		"DATABRICKS_CLIENT_SECRET",
		"DATABRICKS_CLI_PATH",
		"DATABRICKS_AUTH_TYPE",
		"DATABRICKS_DEBUG_TRUNCATE_BYTES",
		"DATABRICKS_DEBUG_HEADERS",
		"DATABRICKS_RATE_LIMIT",
	}

	out := envVars()
	assert.Equal(t, expected, out)
}

func TestAuthProcessEnv(t *testing.T) {
	testutil.ClearEnvironment(t)

	// Environment variables that should be inherited by child processes.
	t.Setenv("HOME", "/home/user")
	t.Setenv("HTTPS_PROXY", "http://proxy.com")

	// Environment variables that should be cleaned up by ProcessEnv():
	t.Setenv("DATABRICKS_HOST", "https://test.com")
	t.Setenv("DATABRICKS_TOKEN", "test-token")
	t.Setenv("DATABRICKS_PASSWORD", "test-password")
	t.Setenv("DATABRICKS_METADATA_SERVICE_URL", "http://somurl.com")
	t.Setenv("ARM_USE_MSI", "true")
	t.Setenv("ARM_TENANT_ID", "test-tenant-id")
	t.Setenv("ARM_CLIENT_ID", "test-client-id")
	t.Setenv("ARM_CLIENT_SECRET", "test-client-secret")

	in := &config.Config{
		Host:  "https://newhost.com",
		Token: "new-token",
	}

	expected := []string{
		"DATABRICKS_HOST=https://newhost.com",
		"DATABRICKS_TOKEN=new-token",
		"HOME=/home/user",
		"HTTPS_PROXY=http://proxy.com",
	}

	out := ProcessEnv(in)
	assert.Equal(t, expected, out)
}