Merge pull request #21 from pedromg/master

Fix on the \n string comparison for the -H flag

client.go

@@ -53,16 +53,16 @@ func StartClient(url_, heads, requestBody string, meth string, dka bool, respons
 	}
 
 	timer := NewTimer()
+
+	hdrs, _ := buildHeaders(heads)
+
 	for {
 		requestBodyReader := strings.NewReader(requestBody)
 		req, _ := http.NewRequest(meth, url_, requestBodyReader)
-		sets := strings.Split(heads, "\n")
 
-		//Split incoming header string by \n and build header pairs
-		for i := range sets {
-			split := strings.SplitN(sets[i], ":", 2)
-			if len(split) == 2 {
-				req.Header.Set(split[0], split[1])
+		for key, vals := range hdrs {
+			for _, val := range vals {
+				req.Header.Set(key, val)
 			}
 		}
 
@@ -102,3 +102,24 @@ func StartClient(url_, heads, requestBody string, meth string, dka bool, respons
 		responseChan <- respObj
 	}
 }
+
+// buildHeaders build the HTTP Request headers from the parsed flag -H or
+// from the default header set.
+// The headers are "set" (not added), thus same key values get replaced.
+// Note: if a key has no value, it is not added into the Headers, by original
+// package design.
+func buildHeaders(heads string) (http.Header, error) {
+
+	heads = strings.Replace(heads, `\n`, "\n", -1)
+	h := http.Header{}
+
+	sets := strings.Split(heads, "\n")
+	for i := range sets {
+		split := strings.SplitN(sets[i], ":", 2)
+		if len(split) == 2 {
+			h.Set(split[0], split[1])
+		}
+	}
+
+	return h, nil
+}

client_test.go

@@ -0,0 +1,70 @@
+package main
+
+import (
+	"net/http"
+	"reflect"
+	"sort"
+	"testing"
+)
+
+func TestBuildHeaders(t *testing.T) {
+
+	var testSet = []struct {
+		H   string // headers passed via flag
+		Res map[string][]string
+	}{
+		{
+			"User-Agent:go-wrk 0.1 bechmark\nContent-Type:text/html;",
+			map[string][]string{
+				"User-Agent":   []string{"go-wrk 0.1 bechmark"},
+				"Content-Type": []string{"text/html;"},
+			},
+		},
+		{
+			"Key:Value",
+			map[string][]string{
+				"Key": []string{"Value"},
+			},
+		},
+		{
+			"Key1:Value1\nKey2:Value2",
+			map[string][]string{
+				"Key1": []string{"Value1"},
+				"Key2": []string{"Value2"},
+			},
+		},
+		{
+			// the headers are set (not added) thus same key values
+			// are replaced.
+			"Key1:Value1A\nKey1:Value1B",
+			map[string][]string{
+				"Key1": []string{"Value1B"},
+			},
+		},
+		{
+			// a key with no value gets removed by design of the package.
+			"Key1",
+			map[string][]string{},
+		},
+	}
+
+	for _, set := range testSet {
+
+		tmpHeaders := http.Header{}
+		for k, v := range set.Res {
+			tmpHeaders[k] = append(tmpHeaders[k], v...)
+			sort.Strings(tmpHeaders[k])
+		}
+
+		headers, _ := buildHeaders(set.H)
+		for _, v := range headers {
+			sort.Strings(v)
+		}
+
+		// comparison; using the not very efficient reflect.DeepEqual
+		// because its a small test suite.
+		if !reflect.DeepEqual(tmpHeaders, headers) {
+			t.Errorf("Different results")
+		}
+	}
+}