added client finalizer
based on the server finazlizer
TRAVIS ALLEN SALAS COX
6 years ago
| 22 | 22 | dec DecodeResponseFunc |
| 23 | 23 | before []RequestFunc |
| 24 | 24 | after []ClientResponseFunc |
| 25 | finalizer ClientFinalizerFunc | |
| 25 | 26 | bufferedStream bool |
| 26 | 27 | } |
| 27 | 28 | |
| 71 | 72 | return func(c *Client) { c.after = append(c.after, after...) } |
| 72 | 73 | } |
| 73 | 74 | |
| 75 | // ClientFinalizer is executed at the end of every HTTP request. | |
| 76 | // By default, no finalizer is registered. | |
| 77 | func ClientFinalizer(f ClientFinalizerFunc) ClientOption { | |
| 78 | return func(s *Client) { s.finalizer = f } | |
| 79 | } | |
| 80 | ||
| 74 | 81 | // BufferedStream sets whether the Response.Body is left open, allowing it |
| 75 | 82 | // to be read from later. Useful for transporting a file as a buffered stream. |
| 76 | 83 | func BufferedStream(buffered bool) ClientOption { |
| 83 | 90 | ctx, cancel := context.WithCancel(ctx) |
| 84 | 91 | defer cancel() |
| 85 | 92 | |
| 86 | req, err := http.NewRequest(c.method, c.tgt.String(), nil) | |
| 93 | // Vars used for client finalizer to ensure there are no nil values | |
| 94 | var ( | |
| 95 | req *http.Request = &http.Request{} | |
| 96 | resp *http.Response = &http.Response{} | |
| 97 | err error | |
| 98 | ) | |
| 99 | if c.finalizer != nil { | |
| 100 | defer func() { | |
| 101 | ctx = context.WithValue(ctx, ContextKeyResponseHeaders, resp.Header) | |
| 102 | ctx = context.WithValue(ctx, ContextKeyResponseSize, resp.ContentLength) | |
| 103 | c.finalizer(ctx, resp.StatusCode, req) | |
| 104 | }() | |
| 105 | } | |
| 106 | ||
| 107 | req, err = http.NewRequest(c.method, c.tgt.String(), nil) | |
| 87 | 108 | if err != nil { |
| 88 | 109 | return nil, err |
| 89 | 110 | } |
| 96 | 117 | ctx = f(ctx, req) |
| 97 | 118 | } |
| 98 | 119 | |
| 99 | resp, err := ctxhttp.Do(ctx, c.client, req) | |
| 120 | resp, err = ctxhttp.Do(ctx, c.client, req) | |
| 100 | 121 | if err != nil { |
| 101 | 122 | return nil, err |
| 102 | 123 | } |
| 124 | ||
| 103 | 125 | if !c.bufferedStream { |
| 104 | 126 | defer resp.Body.Close() |
| 105 | 127 | } |
| 116 | 138 | return response, nil |
| 117 | 139 | } |
| 118 | 140 | } |
| 141 | ||
| 142 | // ClientFinalizerFunc can be used to perform work at the end of a client HTTP | |
| 143 | // request, after the response is returned. The principal | |
| 144 | // intended use is for request logging. In addition to the response code | |
| 145 | // provided in the function signature, additional response parameters are | |
| 146 | // provided in the context under keys with the ContextKeyResponse prefix. | |
| 147 | type ClientFinalizerFunc func(ctx context.Context, code int, r *http.Request) | |
| 119 | 148 | |
| 120 | 149 | // EncodeJSONRequest is an EncodeRequestFunc that serializes the request as a |
| 121 | 150 | // JSON object to the Request body. Many JSON-over-HTTP services can use it as |
| 136 | 136 | } |
| 137 | 137 | if want, have := testbody, string(b); want != have { |
| 138 | 138 | t.Errorf("want %q, have %q", want, have) |
| 139 | } | |
| 140 | } | |
| 141 | ||
| 142 | func TestClientFinalizer(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 | encode = func(context.Context, *http.Request, interface{}) error { return nil } | |
| 150 | decode = func(_ context.Context, r *http.Response) (interface{}, error) { | |
| 151 | return TestResponse{r.Body, ""}, nil | |
| 152 | } | |
| 153 | ) | |
| 154 | ||
| 155 | server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { | |
| 156 | w.Header().Set(headerKey, headerVal) | |
| 157 | w.WriteHeader(statusCode) | |
| 158 | w.Write([]byte(responseBody)) | |
| 159 | })) | |
| 160 | defer server.Close() | |
| 161 | ||
| 162 | client := httptransport.NewClient( | |
| 163 | "GET", | |
| 164 | mustParse(server.URL), | |
| 165 | encode, | |
| 166 | decode, | |
| 167 | httptransport.ClientFinalizer(func(ctx context.Context, code int, _ *http.Request) { | |
| 168 | if want, have := statusCode, code; want != have { | |
| 169 | t.Errorf("StatusCode: want %d, have %d", want, have) | |
| 170 | } | |
| 171 | ||
| 172 | responseHeader := ctx.Value(httptransport.ContextKeyResponseHeaders).(http.Header) | |
| 173 | if want, have := headerVal, responseHeader.Get(headerKey); want != have { | |
| 174 | t.Errorf("%s: want %q, have %q", headerKey, want, have) | |
| 175 | } | |
| 176 | ||
| 177 | responseSize := ctx.Value(httptransport.ContextKeyResponseSize).(int64) | |
| 178 | if want, have := int64(len(responseBody)), responseSize; want != have { | |
| 179 | t.Errorf("response size: want %d, have %d", want, have) | |
| 180 | } | |
| 181 | ||
| 182 | close(done) | |
| 183 | }), | |
| 184 | ) | |
| 185 | ||
| 186 | _, err := client.Endpoint()(context.Background(), struct{}{}) | |
| 187 | if err != nil { | |
| 188 | t.Fatal(err) | |
| 189 | } | |
| 190 | ||
| 191 | select { | |
| 192 | case <-done: | |
| 193 | case <-time.After(time.Second): | |
| 194 | t.Fatal("timeout waiting for finalizer") | |
| 139 | 195 | } |
| 140 | 196 | } |
| 141 | 197 |