Codebase list golang-github-go-kit-kit / 04f878a
Merge pull request #467 from djgilcrease/allow-chaining-before-after Add the ability to specify ServerBefore & ServerAfter multiple times in the ServerOption section Peter Bourgon authored 7 years ago GitHub 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 }