traefik/middlewares/headers_test.go

118 lines
3.2 KiB
Go
Raw Normal View History

2017-06-13 00:48:21 +00:00
package middlewares
2017-11-23 16:40:03 +00:00
// Middleware tests based on https://github.com/unrolled/secure
2017-06-13 00:48:21 +00:00
import (
"net/http"
"net/http/httptest"
"testing"
2017-11-08 10:40:04 +00:00
"github.com/containous/traefik/testhelpers"
"github.com/stretchr/testify/assert"
2017-06-13 00:48:21 +00:00
)
var myHandler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Write([]byte("bar"))
})
2017-11-23 16:40:03 +00:00
// newHeader constructs a new header instance with supplied options.
func newHeader(options ...HeaderOptions) *HeaderStruct {
var o HeaderOptions
if len(options) == 0 {
o = HeaderOptions{}
} else {
o = options[0]
}
return &HeaderStruct{
opt: o,
}
}
2017-06-13 00:48:21 +00:00
func TestNoConfig(t *testing.T) {
2017-11-23 16:40:03 +00:00
header := newHeader()
2017-06-13 00:48:21 +00:00
res := httptest.NewRecorder()
req := testhelpers.MustNewRequest(http.MethodGet, "http://example.com/foo", nil)
2017-11-23 16:40:03 +00:00
header.ServeHTTP(res, req, myHandler)
2017-06-13 00:48:21 +00:00
assert.Equal(t, http.StatusOK, res.Code, "Status not OK")
assert.Equal(t, "bar", res.Body.String(), "Body not the expected")
}
2017-11-23 16:40:03 +00:00
func TestModifyResponseHeaders(t *testing.T) {
header := newHeader(HeaderOptions{
2017-06-13 00:48:21 +00:00
CustomResponseHeaders: map[string]string{
"X-Custom-Response-Header": "test_response",
},
})
res := httptest.NewRecorder()
2017-11-23 16:40:03 +00:00
res.HeaderMap.Add("X-Custom-Response-Header", "test_response")
2017-06-13 00:48:21 +00:00
2017-11-23 16:40:03 +00:00
header.ModifyResponseHeaders(res.Result())
2017-06-13 00:48:21 +00:00
assert.Equal(t, http.StatusOK, res.Code, "Status not OK")
assert.Equal(t, "test_response", res.Header().Get("X-Custom-Response-Header"), "Did not get expected header")
2017-11-23 16:40:03 +00:00
res = httptest.NewRecorder()
res.HeaderMap.Add("X-Custom-Response-Header", "")
header.ModifyResponseHeaders(res.Result())
assert.Equal(t, http.StatusOK, res.Code, "Status not OK")
assert.Equal(t, "", res.Header().Get("X-Custom-Response-Header"), "Did not get expected header")
res = httptest.NewRecorder()
res.HeaderMap.Add("X-Custom-Response-Header", "test_override")
header.ModifyResponseHeaders(res.Result())
assert.Equal(t, http.StatusOK, res.Code, "Status not OK")
assert.Equal(t, "test_override", res.Header().Get("X-Custom-Response-Header"), "Did not get expected header")
2017-06-13 00:48:21 +00:00
}
func TestCustomRequestHeader(t *testing.T) {
2017-11-23 16:40:03 +00:00
header := newHeader(HeaderOptions{
2017-06-13 00:48:21 +00:00
CustomRequestHeaders: map[string]string{
"X-Custom-Request-Header": "test_request",
},
})
res := httptest.NewRecorder()
req := testhelpers.MustNewRequest(http.MethodGet, "/foo", nil)
2017-11-23 16:40:03 +00:00
header.ServeHTTP(res, req, nil)
2017-06-13 00:48:21 +00:00
assert.Equal(t, http.StatusOK, res.Code, "Status not OK")
assert.Equal(t, "test_request", req.Header.Get("X-Custom-Request-Header"), "Did not get expected header")
}
2017-11-23 16:40:03 +00:00
func TestCustomRequestHeaderEmptyValue(t *testing.T) {
header := newHeader(HeaderOptions{
CustomRequestHeaders: map[string]string{
"X-Custom-Request-Header": "test_request",
},
})
res := httptest.NewRecorder()
req := testhelpers.MustNewRequest(http.MethodGet, "/foo", nil)
header.ServeHTTP(res, req, nil)
assert.Equal(t, http.StatusOK, res.Code, "Status not OK")
assert.Equal(t, "test_request", req.Header.Get("X-Custom-Request-Header"), "Did not get expected header")
header = newHeader(HeaderOptions{
CustomRequestHeaders: map[string]string{
"X-Custom-Request-Header": "",
},
})
header.ServeHTTP(res, req, nil)
assert.Equal(t, http.StatusOK, res.Code, "Status not OK")
assert.Equal(t, "", req.Header.Get("X-Custom-Request-Header"), "This header is not expected")
}