traefik/middlewares/error_pages_test.go

package middlewares

import (
	"fmt"
	"net/http"
	"net/http/httptest"
	"strconv"
	"testing"

	"github.com/containous/traefik/types"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"github.com/urfave/negroni"
)

func TestErrorPage(t *testing.T) {
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		fmt.Fprintln(w, "Test Server")
	}))
	defer ts.Close()

	testErrorPage := &types.ErrorPage{Backend: "error", Query: "/test", Status: []string{"500-501", "503-599"}}

	testHandler, err := NewErrorPagesHandler(testErrorPage, ts.URL)
	require.NoError(t, err)

	assert.Equal(t, testHandler.BackendURL, ts.URL+"/test", "Should be equal")

	recorder := httptest.NewRecorder()
	req, err := http.NewRequest(http.MethodGet, ts.URL+"/test", nil)
	require.NoError(t, err)

	handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		fmt.Fprintln(w, "traefik")
	})
	n := negroni.New()
	n.Use(testHandler)
	n.UseHandler(handler)

	n.ServeHTTP(recorder, req)

	assert.Equal(t, http.StatusOK, recorder.Code, "HTTP status")
	assert.Contains(t, recorder.Body.String(), "traefik")

	// ----

	handler500 := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusInternalServerError)
		fmt.Fprintln(w, "oops")
	})
	recorder500 := httptest.NewRecorder()
	n500 := negroni.New()
	n500.Use(testHandler)
	n500.UseHandler(handler500)

	n500.ServeHTTP(recorder500, req)

	assert.Equal(t, http.StatusInternalServerError, recorder500.Code, "HTTP status Internal Server Error")
	assert.Contains(t, recorder500.Body.String(), "Test Server")
	assert.NotContains(t, recorder500.Body.String(), "oops", "Should not return the oops page")

	handler502 := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusBadGateway)
		fmt.Fprintln(w, "oops")
	})
	recorder502 := httptest.NewRecorder()
	n502 := negroni.New()
	n502.Use(testHandler)
	n502.UseHandler(handler502)

	n502.ServeHTTP(recorder502, req)

	assert.Equal(t, http.StatusBadGateway, recorder502.Code, "HTTP status Bad Gateway")
	assert.Contains(t, recorder502.Body.String(), "oops")
	assert.NotContains(t, recorder502.Body.String(), "Test Server", "Should return the oops page since we have not configured the 502 code")
}

func TestErrorPageQuery(t *testing.T) {
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.RequestURI() == "/"+strconv.Itoa(503) {
			fmt.Fprintln(w, "503 Test Server")
		} else {
			fmt.Fprintln(w, "Failed")
		}

	}))
	defer ts.Close()

	testErrorPage := &types.ErrorPage{Backend: "error", Query: "/{status}", Status: []string{"503-503"}}

	testHandler, err := NewErrorPagesHandler(testErrorPage, ts.URL)
	require.NoError(t, err)

	assert.Equal(t, testHandler.BackendURL, ts.URL+"/{status}", "Should be equal")

	handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusServiceUnavailable)
		fmt.Fprintln(w, "oops")
	})

	recorder := httptest.NewRecorder()

	req, err := http.NewRequest(http.MethodGet, ts.URL+"/test", nil)
	require.NoError(t, err)

	n := negroni.New()
	n.Use(testHandler)
	n.UseHandler(handler)

	n.ServeHTTP(recorder, req)

	assert.Equal(t, http.StatusServiceUnavailable, recorder.Code, "HTTP status Service Unavailable")
	assert.Contains(t, recorder.Body.String(), "503 Test Server")
	assert.NotContains(t, recorder.Body.String(), "oops", "Should not return the oops page")
}

func TestErrorPageSingleCode(t *testing.T) {
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.RequestURI() == "/"+strconv.Itoa(503) {
			fmt.Fprintln(w, "503 Test Server")
		} else {
			fmt.Fprintln(w, "Failed")
		}

	}))
	defer ts.Close()

	testErrorPage := &types.ErrorPage{Backend: "error", Query: "/{status}", Status: []string{"503"}}

	testHandler, err := NewErrorPagesHandler(testErrorPage, ts.URL)
	require.NoError(t, err)

	assert.Equal(t, testHandler.BackendURL, ts.URL+"/{status}", "Should be equal")

	handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusServiceUnavailable)
		fmt.Fprintln(w, "oops")
	})

	recorder := httptest.NewRecorder()

	req, err := http.NewRequest(http.MethodGet, ts.URL+"/test", nil)
	require.NoError(t, err)

	n := negroni.New()
	n.Use(testHandler)
	n.UseHandler(handler)

	n.ServeHTTP(recorder, req)

	assert.Equal(t, http.StatusServiceUnavailable, recorder.Code, "HTTP status Service Unavailable")
	assert.Contains(t, recorder.Body.String(), "503 Test Server")
	assert.NotContains(t, recorder.Body.String(), "oops", "Should not return the oops page")
}

func TestNewErrorPagesResponseRecorder(t *testing.T) {
	testCases := []struct {
		desc     string
		rw       http.ResponseWriter
		expected http.ResponseWriter
	}{
		{
			desc:     "Without Close Notify",
			rw:       httptest.NewRecorder(),
			expected: &errorPagesResponseRecorderWithoutCloseNotify{},
		},
		{
			desc:     "With Close Notify",
			rw:       &mockRWCloseNotify{},
			expected: &errorPagesResponseRecorderWithCloseNotify{},
		},
	}

	for _, test := range testCases {
		test := test
		t.Run(test.desc, func(t *testing.T) {
			t.Parallel()

			rec := newErrorPagesResponseRecorder(test.rw)

			assert.IsType(t, rec, test.expected)
		})
	}
}

type mockRWCloseNotify struct{}

func (m *mockRWCloseNotify) CloseNotify() <-chan bool {
	panic("implement me")
}

func (m *mockRWCloseNotify) Header() http.Header {
	panic("implement me")
}

func (m *mockRWCloseNotify) Write([]byte) (int, error) {
	panic("implement me")
}

func (m *mockRWCloseNotify) WriteHeader(int) {
	panic("implement me")
}