From ee5d4f7ceb37a62202fb9943376d01ca0293bc17 Mon Sep 17 00:00:00 2001
From: Panchajanya Mysarla <mpanchajanya@vmware.com>
Date: Tue, 25 Apr 2023 12:33:37 -0500
Subject: [PATCH] Add validators functionality to prompts

---
 component/prompt.go      | 19 +++++++++
 component/prompt_test.go | 83 ++++++++++++++++++++++++++++++++++++++++
 2 files changed, 102 insertions(+)

diff --git a/component/prompt.go b/component/prompt.go
index 59006ad2d..40c4a1ec4 100644
--- a/component/prompt.go
+++ b/component/prompt.go
@@ -112,6 +112,9 @@ type PromptOptions struct {
 	// Standard in/out/error
 	Stdio terminal.Stdio
 	Icons survey.IconSet
+
+	// Validators on user inputs
+	Validators []survey.Validator
 }
 
 // PromptOpt is an option for prompts
@@ -127,6 +130,11 @@ func translatePromptOpts(options *PromptOptions) (surveyOpts []survey.AskOpt) {
 		icons.SelectFocus = options.Icons.SelectFocus
 	}))
 
+	// Add validators
+	for _, v := range options.Validators {
+		surveyOpts = append(surveyOpts, survey.WithValidator(v))
+	}
+
 	return
 }
 
@@ -140,3 +148,14 @@ func WithStdio(in terminal.FileReader, out terminal.FileWriter, err io.Writer) P
 		return nil
 	}
 }
+
+// WithValidator specifies a validator to use while prompting the user
+func WithValidator(v survey.Validator) PromptOpt {
+	return func(options *PromptOptions) error {
+		// add the provided validator to the list
+		options.Validators = append(options.Validators, v)
+
+		// nothing went wrong
+		return nil
+	}
+}
diff --git a/component/prompt_test.go b/component/prompt_test.go
index 8774344ed..2f09a3e4d 100644
--- a/component/prompt_test.go
+++ b/component/prompt_test.go
@@ -4,6 +4,8 @@
 package component
 
 import (
+	"errors"
+	"strings"
 	"testing"
 
 	"github.com/AlecAivazis/survey/v2"
@@ -79,3 +81,84 @@ func Test_PromptOptions(t *testing.T) {
 	assert.Equal("cyan+b", options.Icons.Question.Format)
 	assert.Equal(2, len(opts))
 }
+
+func TestAskPrompt_Validation(t *testing.T) {
+	// Setup mock prompts
+	p := &mockPrompt{
+		answers: []string{"", "very-very-long-name", "ALL_CAPS_NAME", "Test", "test"},
+	}
+
+	var res string
+
+	// Setup PromptOpts validators
+	var promptOpts []PromptOpt
+	noUpperCase := func(v interface{}) error {
+		s := v.(string)
+		if strings.ToLower(s) != s {
+			return errors.New("value contains uppercase characters")
+		}
+		return nil
+	}
+	promptOpts = append(promptOpts, WithValidator(noUpperCase), WithValidator(survey.Required), WithValidator(survey.MaxLength(4)))
+
+	// Prepare the surveyOpts
+	var options = &PromptOptions{}
+	for _, opt := range promptOpts {
+		err := opt(options)
+		assert.Nil(t, err)
+	}
+	surveyOpts := translatePromptOpts(options)
+
+	// Trigger the Prompt
+	err := survey.Ask([]*survey.Question{
+		{
+			Prompt: p,
+		},
+	}, &res, surveyOpts...)
+
+	if err != nil {
+		t.Fatalf("Ask() = %v", err)
+	}
+
+	if res != "test" {
+		t.Errorf("answer: %q, want %q", res, "test")
+	}
+	if p.cleanups != 1 {
+		t.Errorf("cleanups: %d, want %d", p.cleanups, 1)
+	}
+	if err1 := p.printedErrors[0].Error(); err1 != "Value is required" {
+		t.Errorf("printed error 1: %q, want %q", err1, "Value is required")
+	}
+	if err2 := p.printedErrors[1].Error(); err2 != "value is too long. Max length is 4" {
+		t.Errorf("printed error 2: %q, want %q", err2, "value is too long. Max length is 4")
+	}
+	if err3 := p.printedErrors[2].Error(); err3 != "value contains uppercase characters" {
+		t.Errorf("printed error 2: %q, want %q", err3, "value contains uppercase characters")
+	}
+}
+
+type mockPrompt struct {
+	index         int
+	answers       []string
+	cleanups      int
+	printedErrors []error
+}
+
+func (p *mockPrompt) Prompt(*survey.PromptConfig) (interface{}, error) {
+	if p.index >= len(p.answers) {
+		return nil, errors.New("no valid answers provided")
+	}
+	val := p.answers[p.index]
+	p.index++
+	return val, nil
+}
+
+func (p *mockPrompt) Cleanup(*survey.PromptConfig, interface{}) error {
+	p.cleanups++
+	return nil
+}
+
+func (p *mockPrompt) Error(_ *survey.PromptConfig, err error) error {
+	p.printedErrors = append(p.printedErrors, err)
+	return nil
+}