Add testcase to validate the password prompt

Author: apoorvdeshmukh

cmd/sqlcmd/main.go

@@ -102,6 +102,17 @@ func (a SQLCmdArguments) authenticationMethod(hasPassword bool) string {
 	return a.AuthenticationMethod
 }
 
+func getSqlcmdConsole(args SQLCmdArguments) sqlcmd.Console {
+	iactive := args.InputFile == nil && args.Query == ""
+	uactive := args.UserName != ""
+	var sqlcmdConsole sqlcmd.Console = nil
+	if iactive || uactive {
+		sqlcmdConsole = console.NewConsole("")
+		defer sqlcmdConsole.Close()
+	}
+	return sqlcmdConsole
+}
+
 func main() {
 	ctx := kong.Parse(&args, kong.NoDefaultHelp())
 	if args.Help {
@@ -110,9 +121,9 @@ func main() {
 	}
 	vars := sqlcmd.InitializeVariables(!args.DisableCmdAndWarn)
 	setVars(vars, &args)
-
+	line := getSqlcmdConsole(args)
 	// so far sqlcmd prints all the errors itself so ignore it
-	exitCode, _ := run(vars, &args)
+	exitCode, _ := run(vars, &args, line)
 	os.Exit(exitCode)
 }
 
@@ -201,19 +212,13 @@ func setConnect(connect *sqlcmd.ConnectSettings, args *SQLCmdArguments, vars *sq
 	connect.ErrorSeverityLevel = args.ErrorSeverityLevel
 }
 
-func run(vars *sqlcmd.Variables, args *SQLCmdArguments) (int, error) {
+func run(vars *sqlcmd.Variables, args *SQLCmdArguments, line sqlcmd.Console) (int, error) {
 	wd, err := os.Getwd()
 	if err != nil {
 		return 1, err
 	}
 
 	iactive := args.InputFile == nil && args.Query == ""
-	uactive := args.UserName != ""
-	var line sqlcmd.Console = nil
-	if iactive || uactive {
-		line = console.NewConsole("")
-		defer line.Close()
-	}
 
 	s := sqlcmd.New(line, wd, vars)
 	s.UnicodeOutputFile = args.UnicodeOutputFile

cmd/sqlcmd/main_test.go

@@ -138,8 +138,8 @@ func TestRunInputFiles(t *testing.T) {
 	vars := sqlcmd.InitializeVariables(!args.DisableCmdAndWarn)
 	vars.Set(sqlcmd.SQLCMDMAXVARTYPEWIDTH, "0")
 	setVars(vars, &args)
-
-	exitCode, err := run(vars, &args)
+	sqlcmdConsole := getSqlcmdConsole(args)
+	exitCode, err := run(vars, &args, sqlcmdConsole)
 	assert.NoError(t, err, "run")
 	assert.Equal(t, 0, exitCode, "exitCode")
 	bytes, err := os.ReadFile(o.Name())
@@ -162,8 +162,8 @@ func TestUnicodeOutput(t *testing.T) {
 	}
 	vars := sqlcmd.InitializeVariables(!args.DisableCmdAndWarn)
 	setVars(vars, &args)
-
-	exitCode, err := run(vars, &args)
+	sqlcmdConsole := getSqlcmdConsole(args)
+	exitCode, err := run(vars, &args, sqlcmdConsole)
 	assert.NoError(t, err, "run")
 	assert.Equal(t, 0, exitCode, "exitCode")
 	bytes, err := os.ReadFile(o.Name())
@@ -214,7 +214,8 @@ func TestUnicodeInput(t *testing.T) {
 			}
 			vars := sqlcmd.InitializeVariables(!args.DisableCmdAndWarn)
 			setVars(vars, &args)
-			exitCode, err := run(vars, &args)
+			sqlcmdConsole := getSqlcmdConsole(args)
+			exitCode, err := run(vars, &args, sqlcmdConsole)
 			assert.NoError(t, err, "run")
 			assert.Equal(t, 0, exitCode, "exitCode")
 			bytes, err := os.ReadFile(o.Name())
@@ -244,8 +245,8 @@ func TestQueryAndExit(t *testing.T) {
 	vars.Set(sqlcmd.SQLCMDMAXVARTYPEWIDTH, "0")
 	vars.Set("VAR1", "100")
 	setVars(vars, &args)
-
-	exitCode, err := run(vars, &args)
+	sqlcmdConsole := getSqlcmdConsole(args)
+	exitCode, err := run(vars, &args, sqlcmdConsole)
 	assert.NoError(t, err, "run")
 	assert.Equal(t, 0, exitCode, "exitCode")
 	bytes, err := os.ReadFile(o.Name())
@@ -268,17 +269,16 @@ func TestExitOnError(t *testing.T) {
 
 	vars := sqlcmd.InitializeVariables(!args.DisableCmdAndWarn)
 	setVars(vars, &args)
-
-	exitCode, err := run(vars, &args)
+	sqlcmdConsole := getSqlcmdConsole(args)
+	exitCode, err := run(vars, &args, sqlcmdConsole)
 	assert.NoError(t, err, "run")
 	assert.Equal(t, 0, exitCode, "exitCode")
 
 	args.InputFile = []string{"testdata/bad.sql"}
 
 	vars = sqlcmd.InitializeVariables(!args.DisableCmdAndWarn)
 	setVars(vars, &args)
-
-	exitCode, err = run(vars, &args)
+	exitCode, err = run(vars, &args, sqlcmdConsole)
 	assert.NoError(t, err, "run")
 	assert.Equal(t, 1, exitCode, "exitCode")
 
@@ -301,7 +301,8 @@ func TestAzureAuth(t *testing.T) {
 	vars := sqlcmd.InitializeVariables(!args.DisableCmdAndWarn)
 	vars.Set(sqlcmd.SQLCMDMAXVARTYPEWIDTH, "0")
 	setVars(vars, &args)
-	exitCode, err := run(vars, &args)
+	sqlcmdConsole := getSqlcmdConsole(args)
+	exitCode, err := run(vars, &args, sqlcmdConsole)
 	assert.NoError(t, err, "run")
 	assert.Equal(t, 0, exitCode, "exitCode")
 	bytes, err := os.ReadFile(o.Name())
@@ -320,13 +321,65 @@ func TestMissingInputFile(t *testing.T) {
 
 	vars := sqlcmd.InitializeVariables(!args.DisableCmdAndWarn)
 	setVars(vars, &args)
-
-	exitCode, err := run(vars, &args)
+	sqlcmdConsole := getSqlcmdConsole(args)
+	exitCode, err := run(vars, &args, sqlcmdConsole)
 	assert.Error(t, err, "run")
 	assert.Contains(t, err.Error(), "Error occurred while opening or operating on file", "Unexpected error: "+err.Error())
 	assert.Equal(t, 1, exitCode, "exitCode")
 }
 
+func TestPasswordPrompt(t *testing.T) {
+	prompted := false
+	validationConsole := &testConsole{
+		OnPasswordPrompt: func(prompt string) ([]byte, error) {
+			assert.Equal(t, "Password:", prompt, "Incorrect password prompt")
+			prompted = true
+			return []byte{}, nil
+		},
+		OnReadLine: func() (string, error) {
+			assert.Fail(t, "ReadLine should not be called")
+			return "", nil
+		},
+	}
+	args = newArguments()
+	args.InputFile = []string{"testdata/select100.sql"}
+	vars := sqlcmd.InitializeVariables(!args.DisableCmdAndWarn)
+	setVars(vars, &args)
+	exitCode, err := run(vars, &args, validationConsole)
+	assert.False(t, prompted, "Password prompt shown when not expected")
+	assert.NoError(t, err, "run")
+	assert.Equal(t, 0, exitCode, "exitCode")
+
+	args.UserName = "someuser"
+	vars = sqlcmd.InitializeVariables(!args.DisableCmdAndWarn)
+	setVars(vars, &args)
+	exitCode, err = run(vars, &args, validationConsole)
+	assert.True(t, prompted, "Password prompt not displayed for -U")
+	assert.Error(t, err, "run")
+	assert.Equal(t, 1, exitCode, "exitCode")
+}
+
+type testConsole struct {
+	PromptText       string
+	OnPasswordPrompt func(prompt string) ([]byte, error)
+	OnReadLine       func() (string, error)
+}
+
+func (tc *testConsole) Readline() (string, error) {
+	return tc.OnReadLine()
+}
+
+func (tc *testConsole) ReadPassword(prompt string) ([]byte, error) {
+	return tc.OnPasswordPrompt(prompt)
+}
+
+func (tc *testConsole) SetPrompt(s string) {
+	tc.PromptText = s
+}
+
+func (tc *testConsole) Close() {
+}
+
 // Assuming public Azure, use AAD when SQLCMDUSER environment variable is not set
 func canTestAzureAuth() bool {
 	server := os.Getenv(sqlcmd.SQLCMDSERVER)