From 71b1565f2e7a7982859975ef6aff15d5af912fb2 Mon Sep 17 00:00:00 2001
From: Abdullah Gharaibeh <ahg@google.com>
Date: Thu, 16 May 2019 18:21:51 -0400
Subject: [PATCH] Fixed the unit test and added a failure case.

---
 pkg/probe/http/http_test.go | 44 +++++++++++++++++++++++--------------
 1 file changed, 28 insertions(+), 16 deletions(-)

diff --git a/pkg/probe/http/http_test.go b/pkg/probe/http/http_test.go
index 70c4230a393..062fe8f03e9 100644
--- a/pkg/probe/http/http_test.go
+++ b/pkg/probe/http/http_test.go
@@ -316,14 +316,15 @@ func TestHTTPProbeChecker_NonLocalRedirects(t *testing.T) {
 }
 
 func TestHTTPProbeChecker_HostHeaderPreservedAfterRedirect(t *testing.T) {
-	hostHeader := "www.example.com"
+	successHostHeader := "www.success.com"
+	failHostHeader := "www.fail.com"
 
 	handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		switch r.URL.Path {
 		case "/redirect":
 			http.Redirect(w, r, "/success", http.StatusFound)
 		case "/success":
-			if r.Host != hostHeader {
+			if r.Host == successHostHeader {
 				w.WriteHeader(http.StatusOK)
 			} else {
 				http.Error(w, "", http.StatusBadRequest)
@@ -335,18 +336,29 @@ func TestHTTPProbeChecker_HostHeaderPreservedAfterRedirect(t *testing.T) {
 	server := httptest.NewServer(handler)
 	defer server.Close()
 
-	t.Run("local", func(t *testing.T) {
-		prober := New(false)
-		target, err := url.Parse(server.URL + "/redirect")
-		require.NoError(t, err)
-		result, _, _ := prober.Probe(target, nil, wait.ForeverTestTimeout)
-		assert.Equal(t, probe.Success, result)
-	})
-	t.Run("nonlocal", func(t *testing.T) {
-		prober := New(true)
-		target, err := url.Parse(server.URL + "/redirect")
-		require.NoError(t, err)
-		result, _, _ := prober.Probe(target, nil, wait.ForeverTestTimeout)
-		assert.Equal(t, probe.Success, result)
-	})
+	testCases := map[string]struct {
+		hostHeader     string
+		expectedResult probe.Result
+	}{
+		"success": {successHostHeader, probe.Success},
+		"fail":    {failHostHeader, probe.Failure},
+	}
+	for desc, test := range testCases {
+		headers := http.Header{}
+		headers.Add("Host", test.hostHeader)
+		t.Run(desc+"local", func(t *testing.T) {
+			prober := New(false)
+			target, err := url.Parse(server.URL + "/redirect")
+			require.NoError(t, err)
+			result, _, _ := prober.Probe(target, headers, wait.ForeverTestTimeout)
+			assert.Equal(t, test.expectedResult, result)
+		})
+		t.Run(desc+"nonlocal", func(t *testing.T) {
+			prober := New(true)
+			target, err := url.Parse(server.URL + "/redirect")
+			require.NoError(t, err)
+			result, _, _ := prober.Probe(target, headers, wait.ForeverTestTimeout)
+			assert.Equal(t, test.expectedResult, result)
+		})
+	}
 }