Testing for Shorthand Lookup

3 files changed, 51 insertions, 7 deletions

diff --git a/example_test.go b/example_test.go
index bf405f1..4211757 100644
--- a/example_test.go
+++ b/example_test.go
@@ -77,20 +77,26 @@ func Example() {
 }
 
 func ExampleShorthandLookup() {
-	pflag.BoolP("verbose", "v", false, "verbose output")
-
 	name := "verbose"
-	flag := pflag.ShorthandLookup(name[:1])
+	short := name[:1]
+
+	pflag.BoolP(name, short, false, "verbose output")
+
+	// len(short) must be == 1
+	flag := pflag.ShorthandLookup(short)
 
 	fmt.Println(flag.Name)
 }
 
 func ExampleFlagSet_ShorthandLookup() {
+	name := "verbose"
+	short := name[:1]
+
 	fs := pflag.NewFlagSet("Example", pflag.ContinueOnError)
-	fs.BoolP("verbose", "v", false, "verbose output")
+	fs.BoolP(name, short, false, "verbose output")
 
-	name := "verbose"
-	flag := fs.ShorthandLookup(name[:1])
+	// len(short) must be == 1
+	flag := fs.ShorthandLookup(short)
 
 	fmt.Println(flag.Name)
 }
diff --git a/flag.go b/flag.go
index 0519465..6f1fc30 100644
--- a/flag.go
+++ b/flag.go
@@ -327,7 +327,9 @@ func (f *FlagSet) ShorthandLookup(name string) *Flag {
 		return nil
 	}
 	if len(name) > 1 {
-		panic("can't look up for a shorthand with name more than one character")
+		msg := fmt.Sprintf("can not look up shorthand which is more than one ASCII character: %q", name)
+		fmt.Fprintf(f.out(), msg)
+		panic(msg)
 	}
 	c := name[0]
 	return f.shorthands[c]
diff --git a/flag_test.go b/flag_test.go
index 8248277..c3def0f 100644
--- a/flag_test.go
+++ b/flag_test.go
@@ -446,6 +446,42 @@ func TestShorthand(t *testing.T) {
 	}
 }
 
+func TestShorthandLookup(t *testing.T) {
+	f := NewFlagSet("shorthand", ContinueOnError)
+	if f.Parsed() {
+		t.Error("f.Parse() = true before Parse")
+	}
+	f.BoolP("boola", "a", false, "bool value")
+	f.BoolP("boolb", "b", false, "bool2 value")
+	args := []string{
+		"-ab",
+	}
+	f.SetOutput(ioutil.Discard)
+	if err := f.Parse(args); err != nil {
+		t.Error("expected no error, got ", err)
+	}
+	if !f.Parsed() {
+		t.Error("f.Parse() = false after Parse")
+	}
+	flag := f.ShorthandLookup("a")
+	if flag == nil {
+		t.Errorf("f.ShorthandLookup(\"a\") returned nil")
+	}
+	if flag.Name != "boola" {
+		t.Errorf("f.ShorthandLookup(\"a\") found %q instead of \"boola\"", flag.Name)
+	}
+	flag = f.ShorthandLookup("")
+	if flag != nil {
+		t.Errorf("f.ShorthandLookup(\"\") did not return nil")
+	}
+	defer func() {
+		recover()
+	}()
+	flag = f.ShorthandLookup("ab")
+	// should NEVER get here. lookup should panic. defer'd func should recover it.
+	t.Errorf("f.ShorthandLookup(\"ab\") did not panic")
+}
+
 func TestParse(t *testing.T) {
 	ResetForTesting(func() { t.Error("bad parse") })
 	testParse(GetCommandLine(), t)