Move the interactive device code login tests to the terminal emulator

cmd/auth/login_device_code_test.go

@@ -3,20 +3,27 @@ package auth
 import (
 	"bytes"
 	"context"
+	"encoding/json"
 	"errors"
 	"io"
+	"net/url"
+	"os"
+	"path/filepath"
 	"runtime"
+	"strconv"
 	"strings"
 	"testing"
 	"time"
 
+	"github.com/databricks/cli/internal/testutil"
 	"github.com/databricks/cli/libs/auth"
 	"github.com/databricks/cli/libs/auth/u2m"
 	"github.com/databricks/cli/libs/cmdio"
 	"github.com/databricks/cli/libs/cmdio/cmdiotest/termtest"
 	"github.com/databricks/cli/libs/databrickscfg/profile"
 	"github.com/databricks/cli/libs/env"
 	"github.com/databricks/cli/libs/flags"
+	"github.com/databricks/cli/libs/testserver"
 	"github.com/spf13/cobra"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -211,6 +218,7 @@ func TestIsDeviceCodeSupportedForHost(t *testing.T) {
 // startDeviceCodeLogin runs "auth login --use-device-code" with args in a
 // terminal emulator, using config as the contents of .databrickscfg.
 func startDeviceCodeLogin(t *testing.T, config string, args ...string) *termtest.Term[struct{}] {
+	testutil.CleanupEnvironment(t)
 	home := t.TempDir()
 	t.Setenv("HOME", home)
 	t.Setenv("USERPROFILE", home)
@@ -219,7 +227,9 @@ func startDeviceCodeLogin(t *testing.T, config string, args ...string) *termtest
 	t.Setenv(deviceCodeEnvVar, "true")
 
 	return termtest.Run(t, func(ctx context.Context) (struct{}, error) {
-		cmd := newLoginCommand(&auth.AuthArguments{})
+		authArguments := &auth.AuthArguments{}
+		cmd := newLoginCommand(authArguments)
+		cmd.Flags().StringVar(&authArguments.Host, "host", "", "")
 		cmd.Flags().String("profile", "", "")
 		cmd.SetContext(ctx)
 		cmd.SetArgs(append([]string{"--use-device-code"}, args...))
@@ -304,3 +314,86 @@ resources = https://workspace.test/ai-gateway/mcp-services/system.ai.github
 	_, err := tm.Result()
 	assert.EqualError(t, err, "^C")
 }
+
+// newDeviceCodeServer returns a fake workspace with device code login.
+func newDeviceCodeServer(t *testing.T) *testserver.Server {
+	server := testserver.New(t)
+	testserver.AddDefaultHandlers(server)
+	return server
+}
+
+// A wrong or blank confirmation code leaves the device code valid, so the user
+// can enter it again. The code is accepted the way the browser displays it, in
+// any case.
+func TestLoginDeviceCodeReadsConfirmationCode(t *testing.T) {
+	server := newDeviceCodeServer(t)
+	tm := startDeviceCodeLogin(t, "[test]\nhost = "+server.URL+"\n", "--profile", "test")
+	tm.WaitFor("Confirmation code")
+	tm.Type("AAAA-AAAA\n")
+	tm.WaitFor("Enter the confirmation code again (2 attempts left)")
+	tm.Type("\n")
+	tm.WaitFor("Enter the confirmation code again (1 attempt left)")
+	tm.Type("ftfa-qqky\n")
+	_, err := tm.Result()
+	require.NoError(t, err)
+
+	p, err := loadProfileByName(t.Context(), "test", profile.DefaultProfiler)
+	require.NoError(t, err)
+	require.NotNil(t, p)
+	assert.Equal(t, server.URL, p.Host)
+	assert.Equal(t, strconv.Itoa(testserver.TestWorkspaceID), p.WorkspaceID)
+	assert.Equal(t, authTypeDatabricksCLI, p.AuthType)
+
+	raw, err := os.ReadFile(filepath.Join(os.Getenv("HOME"), ".databricks", "token-cache.json"))
+	require.NoError(t, err)
+	var cache struct {
+		Tokens map[string]json.RawMessage `json:"tokens"`
+	}
+	require.NoError(t, json.Unmarshal(raw, &cache))
+	assert.Contains(t, cache.Tokens, "test")
+}
+
+func TestLoginDeviceCodeFailsAfterWrongConfirmationCodes(t *testing.T) {
+	server := newDeviceCodeServer(t)
+	tm := startDeviceCodeLogin(t, "[test]\nhost = "+server.URL+"\n", "--profile", "test")
+	tm.WaitFor("Confirmation code")
+	tm.Type("AAAA-AAAA\nBBBB-BBBB\nCCCC-CCCC\n")
+	_, err := tm.Result()
+	assert.EqualError(t, err, "Invalid confirmation code (error code: invalid_request). Start login from scratch")
+}
+
+func TestLoginDeviceCodeFailsWithoutConfirmationCode(t *testing.T) {
+	server := newDeviceCodeServer(t)
+	tm := startDeviceCodeLogin(t, "[test]\nhost = "+server.URL+"\n", "--profile", "test")
+	tm.WaitFor("Confirmation code")
+	tm.Close()
+	_, err := tm.Result()
+	assert.EqualError(t, err, "reading confirmation code: EOF")
+}
+
+// The token gets the scopes requested at device authorization, so --scopes is
+// sent there and not with the token request.
+func TestLoginDeviceCodeRequestsScopes(t *testing.T) {
+	server := newDeviceCodeServer(t)
+	var scope string
+	server.RequestCallback = func(r *testserver.Request) {
+		if r.URL.Path == "/oidc/v1/device/auth" {
+			form, err := url.ParseQuery(string(r.Body))
+			assert.NoError(t, err)
+			scope = form.Get("scope")
+		}
+	}
+	tm := startDeviceCodeLogin(t, "", "--host", server.URL, "--profile", "test", "--scopes", "sql")
+	tm.WaitFor("Confirmation code")
+	tm.Type("FTFA-QQKY\n")
+	_, err := tm.Result()
+	require.NoError(t, err)
+	assert.Equal(t, "offline_access sql", scope)
+
+	p, err := loadProfileByName(t.Context(), "test", profile.DefaultProfiler)
+	require.NoError(t, err)
+	require.NotNil(t, p)
+	assert.Equal(t, server.URL, p.Host)
+	assert.Equal(t, strconv.Itoa(testserver.TestWorkspaceID), p.WorkspaceID)
+	assert.Equal(t, "sql", p.Scopes)
+}