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
58 | 58 | // ServerBefore functions are executed on the HTTP request object before the |
59 | 59 | // request is decoded. |
60 | 60 | func ServerBefore(before ...RequestFunc) ServerOption { |
61 | return func(s *Server) { s.before = before } | |
61 | return func(s *Server) { s.before = append(s.before, before...) } | |
62 | 62 | } |
63 | 63 | |
64 | 64 | // ServerAfter functions are executed on the HTTP response writer after the |
65 | 65 | // endpoint is invoked, but before anything is written to the client. |
66 | 66 | func ServerAfter(after ...ResponseFunc) ServerOption { |
67 | return func(s *Server) { s.after = after } | |
67 | return func(s *Server) { s.after = append(s.after, after...) } | |
68 | 68 | } |
69 | 69 | |
70 | 70 | // ServerErrorLogger is used to log non-terminal errors. By default, no errors |
50 | 50 | // ServerBefore functions are executed on the HTTP request object before the |
51 | 51 | // request is decoded. |
52 | 52 | func ServerBefore(before ...RequestFunc) ServerOption { |
53 | return func(s *Server) { s.before = before } | |
53 | return func(s *Server) { s.before = append(s.before, before...) } | |
54 | 54 | } |
55 | 55 | |
56 | 56 | // ServerAfter functions are executed on the HTTP response writer after the |
57 | 57 | // endpoint is invoked, but before anything is written to the client. |
58 | 58 | func ServerAfter(after ...ServerResponseFunc) ServerOption { |
59 | return func(s *Server) { s.after = after } | |
59 | return func(s *Server) { s.after = append(s.after, after...) } | |
60 | 60 | } |
61 | 61 | |
62 | 62 | // ServerErrorEncoder is used to encode errors to the http.ResponseWriter |
89 | 89 | buf, _ := ioutil.ReadAll(resp.Body) |
90 | 90 | if want, have := http.StatusOK, resp.StatusCode; want != have { |
91 | 91 | 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") | |
92 | 185 | } |
93 | 186 | } |
94 | 187 |
44 | 44 | // ServerBefore functions are executed on the HTTP request object before the |
45 | 45 | // request is decoded. |
46 | 46 | func ServerBefore(before ...RequestFunc) ServerOption { |
47 | return func(s *Server) { s.before = before } | |
47 | return func(s *Server) { s.before = append(s.before, before...) } | |
48 | 48 | } |
49 | 49 | |
50 | 50 | // ServeHTTP implements http.Handler. |
118 | 118 | t.Errorf("want %d or %d, have %d", http.StatusBadGateway, http.StatusInternalServerError, resp.StatusCode) |
119 | 119 | } |
120 | 120 | } |
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 | } |