Function TestModifyResponseUseContext

middleware/proxy_test.go:905–939

(t *testing.T)

}

func TestModifyResponseUseContext(t *testing.T) {
	server := httptest.NewServer(
		http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			w.WriteHeader(http.StatusOK)
			w.Write([]byte("OK"))
		}),
	)
	defer server.Close()
	targetURL, _ := url.Parse(server.URL)
	e := echo.New()
	e.Use(ProxyWithConfig(
		ProxyConfig{
			Balancer: &customBalancer{
				target: &ProxyTarget{
					Name: "tst",
					URL:  targetURL,
				},
			},
			RetryCount: 1,
			ModifyResponse: func(res *http.Response) error {
				val := res.Request.Context().Value(testContextKey("FROM_BALANCER"))
				if valStr, ok := val.(string); ok {
					res.Header.Set("FROM_BALANCER", valStr)
				}
				return nil
			},
		},
	))
	req := httptest.NewRequest(http.MethodGet, "/", nil)
	rec := httptest.NewRecorder()
	e.ServeHTTP(rec, req)
	assert.Equal(t, http.StatusOK, rec.Code)
	assert.Equal(t, "OK", rec.Body.String())
	assert.Equal(t, "CUSTOM_BALANCER", rec.Header().Get("FROM_BALANCER"))
}

func createSimpleWebSocketServer(serveTLS bool) *httptest.Server {
	handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {