Codebase list golang-github-go-kit-kit / c647984
Add the ability to specify ServerBefore & ServerAfter multiple times in the ServerOption section This makes it so you can set some global ServerBefore or After options without having to make sure you specify them individually on every handler Dj Gilcrease authored 7 years ago Dj Gilcrease committed 7 years ago
5 changed file(s) with 140 addition(s) and 5 deletion(s). Raw diff Collapse all Expand all
5858 // ServerBefore functions are executed on the HTTP request object before the
5959 // request is decoded.
6060 func ServerBefore(before ...RequestFunc) ServerOption {
61 return func(s *Server) { s.before = before }
61 return func(s *Server) { s.before = append(s.before, before...) }
6262 }
6363
6464 // ServerAfter functions are executed on the HTTP response writer after the
6565 // endpoint is invoked, but before anything is written to the client.
6666 func ServerAfter(after ...ResponseFunc) ServerOption {
67 return func(s *Server) { s.after = after }
67 return func(s *Server) { s.after = append(s.after, after...) }
6868 }
6969
7070 // ServerErrorLogger is used to log non-terminal errors. By default, no errors
5050 // ServerBefore functions are executed on the HTTP request object before the
5151 // request is decoded.
5252 func ServerBefore(before ...RequestFunc) ServerOption {
53 return func(s *Server) { s.before = before }
53 return func(s *Server) { s.before = append(s.before, before...) }
5454 }
5555
5656 // ServerAfter functions are executed on the HTTP response writer after the
5757 // endpoint is invoked, but before anything is written to the client.
5858 func ServerAfter(after ...ServerResponseFunc) ServerOption {
59 return func(s *Server) { s.after = after }
59 return func(s *Server) { s.after = append(s.after, after...) }
6060 }
6161
6262 // ServerErrorEncoder is used to encode errors to the http.ResponseWriter
8989 buf, _ := ioutil.ReadAll(resp.Body)
9090 if want, have := http.StatusOK, resp.StatusCode; want != have {
9191 t.Errorf("want %d, have %d (%s)", want, have, buf)
92 }
93 }
94
95
96 func TestMultipleServerBefore(t *testing.T) {
97 var (
98 headerKey = "X-Henlo-Lizer"
99 headerVal = "Helllo you stinky lizard"
100 statusCode = http.StatusTeapot
101 responseBody = "go eat a fly ugly\n"
102 done = make(chan struct{})
103 )
104 handler := httptransport.NewServer(
105 context.Background(),
106 endpoint.Nop,
107 func(context.Context, *http.Request) (interface{}, error) {
108 return struct{}{}, nil
109 },
110 func(_ context.Context, w http.ResponseWriter, _ interface{}) error {
111 w.Header().Set(headerKey, headerVal)
112 w.WriteHeader(statusCode)
113 w.Write([]byte(responseBody))
114 return nil
115 },
116 httptransport.ServerBefore(func(ctx context.Context, r *http.Request) context.Context {
117 ctx = context.WithValue(ctx, "one", 1)
118
119 return ctx
120 }),
121 httptransport.ServerBefore(func(ctx context.Context, r *http.Request) context.Context {
122 if _, ok := ctx.Value("one").(int); !ok {
123 t.Error("Value was not set properly when multiple ServerBefores are used")
124 }
125
126 close(done)
127 return ctx
128 }),
129 )
130
131 server := httptest.NewServer(handler)
132 defer server.Close()
133 go http.Get(server.URL)
134
135 select {
136 case <-done:
137 case <-time.After(time.Second):
138 t.Fatal("timeout waiting for finalizer")
139 }
140 }
141
142 func TestMultipleServerAfter(t *testing.T) {
143 var (
144 headerKey = "X-Henlo-Lizer"
145 headerVal = "Helllo you stinky lizard"
146 statusCode = http.StatusTeapot
147 responseBody = "go eat a fly ugly\n"
148 done = make(chan struct{})
149 )
150 handler := httptransport.NewServer(
151 context.Background(),
152 endpoint.Nop,
153 func(context.Context, *http.Request) (interface{}, error) {
154 return struct{}{}, nil
155 },
156 func(_ context.Context, w http.ResponseWriter, _ interface{}) error {
157 w.Header().Set(headerKey, headerVal)
158 w.WriteHeader(statusCode)
159 w.Write([]byte(responseBody))
160 return nil
161 },
162 httptransport.ServerAfter(func(ctx context.Context, w http.ResponseWriter) context.Context {
163 ctx = context.WithValue(ctx, "one", 1)
164
165 return ctx
166 }),
167 httptransport.ServerAfter(func(ctx context.Context, w http.ResponseWriter) context.Context {
168 if _, ok := ctx.Value("one").(int); !ok {
169 t.Error("Value was not set properly when multiple ServerAfters are used")
170 }
171
172 close(done)
173 return ctx
174 }),
175 )
176
177 server := httptest.NewServer(handler)
178 defer server.Close()
179 go http.Get(server.URL)
180
181 select {
182 case <-done:
183 case <-time.After(time.Second):
184 t.Fatal("timeout waiting for finalizer")
92185 }
93186 }
94187
4444 // ServerBefore functions are executed on the HTTP request object before the
4545 // request is decoded.
4646 func ServerBefore(before ...RequestFunc) ServerOption {
47 return func(s *Server) { s.before = before }
47 return func(s *Server) { s.before = append(s.before, before...) }
4848 }
4949
5050 // ServeHTTP implements http.Handler.
118118 t.Errorf("want %d or %d, have %d", http.StatusBadGateway, http.StatusInternalServerError, resp.StatusCode)
119119 }
120120 }
121
122 func TestMultipleServerBefore(t *testing.T) {
123 const (
124 headerKey = "X-TEST-HEADER"
125 headerVal = "go-kit-proxy"
126 )
127
128 originServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
129 if want, have := headerVal, r.Header.Get(headerKey); want != have {
130 t.Errorf("want %q, have %q", want, have)
131 }
132
133 w.WriteHeader(http.StatusOK)
134 w.Write([]byte("hey"))
135 }))
136 defer originServer.Close()
137 originURL, _ := url.Parse(originServer.URL)
138
139 handler := httptransport.NewServer(
140 context.Background(),
141 originURL,
142 httptransport.ServerBefore(func(ctx context.Context, r *http.Request) context.Context {
143 r.Header.Add(headerKey, headerVal)
144 return ctx
145 }),
146 httptransport.ServerBefore(func(ctx context.Context, r *http.Request) context.Context {
147 return ctx
148 }),
149 )
150 proxyServer := httptest.NewServer(handler)
151 defer proxyServer.Close()
152
153 resp, _ := http.Get(proxyServer.URL)
154 if want, have := http.StatusOK, resp.StatusCode; want != have {
155 t.Errorf("want %d, have %d", want, have)
156 }
157
158 responseBody, _ := ioutil.ReadAll(resp.Body)
159 if want, have := "hey", string(responseBody); want != have {
160 t.Errorf("want %q, have %q", want, have)
161 }
162 }