Increased unit test coverage for pkg/protocol/header_test (#1074)

pkg/protocol/header_test.go

@@ -69,18 +69,18 @@ func TestResponseHeaderSetHeaderLength(t *testing.T) {
 
 func TestSetNoHTTP11(t *testing.T) {
 	rh := ResponseHeader{}
-	rh.SetNoHTTP11(true)
+	rh.SetProtocol(consts.HTTP10)
 	assert.DeepEqual(t, consts.HTTP10, rh.protocol)
 
-	rh.SetNoHTTP11(false)
+	rh.SetProtocol(consts.HTTP11)
 	assert.DeepEqual(t, consts.HTTP11, rh.protocol)
 	assert.True(t, rh.IsHTTP11())
 
 	h := RequestHeader{}
-	h.SetNoHTTP11(true)
+	h.SetProtocol(consts.HTTP10)
 	assert.DeepEqual(t, consts.HTTP10, h.protocol)
 
-	h.SetNoHTTP11(false)
+	h.SetProtocol(consts.HTTP11)
 	assert.DeepEqual(t, consts.HTTP11, h.protocol)
 	assert.True(t, h.IsHTTP11())
 }
@@ -101,6 +101,17 @@ func TestSetContentLengthBytes(t *testing.T) {
 	assert.DeepEqual(t, rh.contentLengthBytes, []byte("foo"))
 }
 
+func TestInitContentLengthWithValue(t *testing.T) {
+	initLength := 100
+	h := RequestHeader{}
+	h.InitContentLengthWithValue(initLength)
+	assert.DeepEqual(t, h.contentLength, initLength)
+
+	rh := ResponseHeader{}
+	rh.InitContentLengthWithValue(initLength)
+	assert.DeepEqual(t, rh.contentLength, initLength)
+}
+
 func TestSetContentEncoding(t *testing.T) {
 	rh := ResponseHeader{}
 	rh.SetContentEncoding("gzip")
@@ -171,24 +182,71 @@ func TestResponseHeaderGet(t *testing.T) {
 	assert.DeepEqual(t, val, rightVal)
 }
 
+func TestRequestHeaderGetAll(t *testing.T) {
+	h := RequestHeader{}
+	h.Set("Foo-Bar", "foo")
+	h.Add("Foo-Bar", "bar")
+	h.Add("Foo-Bar", "foo-bar")
+	values := h.GetAll("Foo-Bar")
+	assert.DeepEqual(t, values, []string{"foo", "bar", "foo-bar"})
+}
+
+func TestResponseHeaderGetAll(t *testing.T) {
+	h := ResponseHeader{}
+	h.Set("Foo-Bar", "foo")
+	h.Add("Foo-Bar", "bar")
+	h.Add("Foo-Bar", "foo-bar")
+	values := h.GetAll("Foo-Bar")
+	assert.DeepEqual(t, values, []string{"foo", "bar", "foo-bar"})
+}
+
 func TestRequestHeaderVisitAll(t *testing.T) {
 	h := RequestHeader{}
 	h.Set("xxx", "yyy")
 	h.Set("xxx2", "yyy2")
-	h.VisitAll(
-		func(k, v []byte) {
-			key := string(k)
-			value := string(v)
-			if key != "Xxx" && key != "Xxx2" {
-				t.Fatalf("Unexpected %v. Expected %v", key, "xxx or yyy")
-			}
-			if key == "Xxx" && value != "yyy" {
-				t.Fatalf("Unexpected %v. Expected %v", value, "yyy")
-			}
-			if key == "Xxx2" && value != "yyy2" {
-				t.Fatalf("Unexpected %v. Expected %v", value, "yyy2")
-			}
-		})
+	h.SetHost("host")
+	h.SetContentLengthBytes([]byte("content-length"))
+	h.Set(consts.HeaderContentType, "content-type")
+	h.Set(consts.HeaderUserAgent, "user-agent")
+	err := h.Trailer().SetTrailers([]byte("foo, bar"))
+	if err != nil {
+		t.Fatalf("Set trailer err %v", err)
+	}
+	h.SetCookie("foo", "bar")
+	h.Set(consts.HeaderConnection, "close")
+	h.VisitAll(func(k, v []byte) {
+		key := string(k)
+		value := string(v)
+		switch key {
+		case consts.HeaderHost:
+			assert.DeepEqual(t, value, "host")
+		case consts.HeaderContentLength:
+			assert.DeepEqual(t, value, "content-length")
+		case consts.HeaderContentType:
+			assert.DeepEqual(t, value, "content-type")
+		case consts.HeaderUserAgent:
+			assert.DeepEqual(t, value, "user-agent")
+		case consts.HeaderTrailer:
+			assert.DeepEqual(t, value, "Foo, Bar")
+		case consts.HeaderCookie:
+			assert.DeepEqual(t, value, "foo=bar")
+		case consts.HeaderConnection:
+			assert.DeepEqual(t, value, "close")
+		case "Xxx":
+			assert.DeepEqual(t, value, "yyy")
+		case "Xxx2":
+			assert.DeepEqual(t, value, "yyy2")
+		default:
+			t.Fatalf("Unexpected key %v", key)
+		}
+	})
+}
+
+func TestRequestHeaderCookie(t *testing.T) {
+	var h RequestHeader
+	h.SetCookie("foo", "bar")
+	cookie := h.Cookie("foo")
+	assert.DeepEqual(t, []byte("bar"), cookie)
 }
 
 func TestRequestHeaderCookies(t *testing.T) {
@@ -215,6 +273,7 @@ func TestRequestHeaderDel(t *testing.T) {
 	h.Set(consts.HeaderServer, "aaabbb")
 	h.Set(consts.HeaderContentLength, "1123")
 	h.Set(consts.HeaderTrailer, "foo, bar")
+	h.Set(consts.HeaderUserAgent, "foo-bar")
 	h.SetHost("foobar")
 	h.SetCookie("foo", "bar")
 
@@ -226,6 +285,7 @@ func TestRequestHeaderDel(t *testing.T) {
 	h.del([]byte("Set-Cookie"))
 	h.del([]byte("Host"))
 	h.del([]byte(consts.HeaderTrailer))
+	h.del([]byte(consts.HeaderUserAgent))
 	h.DelCookie("foo")
 	h.Del("ccc")
 
@@ -269,6 +329,10 @@ func TestRequestHeaderDel(t *testing.T) {
 	if len(hv) > 0 {
 		t.Fatalf("non-zero value: %q", hv)
 	}
+	hv = h.Peek(consts.HeaderUserAgent)
+	if len(hv) > 0 {
+		t.Fatalf("non-zero value: %q", hv)
+	}
 	if h.ContentLength() != 0 {
 		t.Fatalf("unexpected content-length: %d. Expecting 0", h.ContentLength())
 	}
@@ -584,6 +648,16 @@ func TestRequestHeaderDelAllCookies(t *testing.T) {
 	}
 }
 
+func TestResponseHeaderDelAllCookies(t *testing.T) {
+	var h ResponseHeader
+	h.SetCanonical([]byte(consts.HeaderSetCookie), []byte("foo"))
+	h.DelAllCookies()
+	hv := h.FullCookie()
+	if len(hv) > 0 {
+		t.Fatalf("non-zero value: %q", hv)
+	}
+}
+
 func TestRequestHeaderSetNoDefaultContentType(t *testing.T) {
 	var h RequestHeader
 	h.SetMethod(http.MethodPost)