Fixed the unit test and added a failure case.

This commit is contained in:
Abdullah Gharaibeh 2019-05-16 18:21:51 -04:00
parent 792ab9d3d0
commit 71b1565f2e

View File

@ -316,14 +316,15 @@ func TestHTTPProbeChecker_NonLocalRedirects(t *testing.T) {
} }
func TestHTTPProbeChecker_HostHeaderPreservedAfterRedirect(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) { handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
switch r.URL.Path { switch r.URL.Path {
case "/redirect": case "/redirect":
http.Redirect(w, r, "/success", http.StatusFound) http.Redirect(w, r, "/success", http.StatusFound)
case "/success": case "/success":
if r.Host != hostHeader { if r.Host == successHostHeader {
w.WriteHeader(http.StatusOK) w.WriteHeader(http.StatusOK)
} else { } else {
http.Error(w, "", http.StatusBadRequest) http.Error(w, "", http.StatusBadRequest)
@ -335,18 +336,29 @@ func TestHTTPProbeChecker_HostHeaderPreservedAfterRedirect(t *testing.T) {
server := httptest.NewServer(handler) server := httptest.NewServer(handler)
defer server.Close() defer server.Close()
t.Run("local", func(t *testing.T) { 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) prober := New(false)
target, err := url.Parse(server.URL + "/redirect") target, err := url.Parse(server.URL + "/redirect")
require.NoError(t, err) require.NoError(t, err)
result, _, _ := prober.Probe(target, nil, wait.ForeverTestTimeout) result, _, _ := prober.Probe(target, headers, wait.ForeverTestTimeout)
assert.Equal(t, probe.Success, result) assert.Equal(t, test.expectedResult, result)
}) })
t.Run("nonlocal", func(t *testing.T) { t.Run(desc+"nonlocal", func(t *testing.T) {
prober := New(true) prober := New(true)
target, err := url.Parse(server.URL + "/redirect") target, err := url.Parse(server.URL + "/redirect")
require.NoError(t, err) require.NoError(t, err)
result, _, _ := prober.Probe(target, nil, wait.ForeverTestTimeout) result, _, _ := prober.Probe(target, headers, wait.ForeverTestTimeout)
assert.Equal(t, probe.Success, result) assert.Equal(t, test.expectedResult, result)
}) })
}
} }