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 {
prober := New(false) hostHeader string
target, err := url.Parse(server.URL + "/redirect") expectedResult probe.Result
require.NoError(t, err) }{
result, _, _ := prober.Probe(target, nil, wait.ForeverTestTimeout) "success": {successHostHeader, probe.Success},
assert.Equal(t, probe.Success, result) "fail": {failHostHeader, probe.Failure},
}) }
t.Run("nonlocal", func(t *testing.T) { for desc, test := range testCases {
prober := New(true) headers := http.Header{}
target, err := url.Parse(server.URL + "/redirect") headers.Add("Host", test.hostHeader)
require.NoError(t, err) t.Run(desc+"local", func(t *testing.T) {
result, _, _ := prober.Probe(target, nil, wait.ForeverTestTimeout) prober := New(false)
assert.Equal(t, probe.Success, result) 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)
})
}
} }