Merge pull request #29 from liggitt/set_annotation

Add SetAnnotation helper

1 files changed, 32 insertions, 0 deletions

diff --git a/flag_test.go b/flag_test.go
index d3c1714..40140d2 100644
--- a/flag_test.go
+++ b/flag_test.go
@@ -11,6 +11,7 @@ import (
 	"io/ioutil"
 	"net"
 	"os"
+	"reflect"
 	"sort"
 	"strings"
 	"testing"
@@ -107,6 +108,37 @@ func TestUsage(t *testing.T) {
 	}
 }
 
+func TestAnnotation(t *testing.T) {
+	f := NewFlagSet("shorthand", ContinueOnError)
+
+	if err := f.SetAnnotation("missing-flag", "key", nil); err == nil {
+		t.Errorf("Expected error setting annotation on non-existent flag")
+	}
+
+	f.StringP("stringa", "a", "", "string value")
+	if err := f.SetAnnotation("stringa", "key", nil); err != nil {
+		t.Errorf("Unexpected error setting new nil annotation: %v", err)
+	}
+	if annotation := f.Lookup("stringa").Annotations["key"]; annotation != nil {
+		t.Errorf("Unexpected annotation: %v", annotation)
+	}
+
+	f.StringP("stringb", "b", "", "string2 value")
+	if err := f.SetAnnotation("stringb", "key", []string{"value1"}); err != nil {
+		t.Errorf("Unexpected error setting new annotation: %v", err)
+	}
+	if annotation := f.Lookup("stringb").Annotations["key"]; !reflect.DeepEqual(annotation, []string{"value1"}) {
+		t.Errorf("Unexpected annotation: %v", annotation)
+	}
+
+	if err := f.SetAnnotation("stringb", "key", []string{"value2"}); err != nil {
+		t.Errorf("Unexpected error updating annotation: %v", err)
+	}
+	if annotation := f.Lookup("stringb").Annotations["key"]; !reflect.DeepEqual(annotation, []string{"value2"}) {
+		t.Errorf("Unexpected annotation: %v", annotation)
+	}
+}
+
 func testParse(f *FlagSet, t *testing.T) {
 	if f.Parsed() {
 		t.Error("f.Parse() = true before Parse")