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
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 | } |