如何在ReverseProxy中测试标题？

Question

I am trying to unit test the following code:

func(h *Handler)Forward(w http.ResponseWriter, r *http.Request) {

    url, err := url.Parse("http://test.com")
    if err != nil {
       return
    }

    reverseProxy := &httputil.ReverseProxy{
        Director: func(r *http.Request) {
            r.URL.Host = url.Host
            r.URL.Path = "/"
            r.URL.Scheme = url.Scheme
            r.Host = url.Host
            r.Header.Set("X-Forwarded-Host", r.Header.Get("Host"))
        },
    }


    reverseProxy.ServeHTTP(w, r)
}

I am not able to figure out how to test whether headers are being modified by the Director function. How do we test headers in a reverseproxy in Go?

Answer 1

1. Inject external dependencies into your unit under test

The biggest problem I can see right now is that the URL you forward to is hard-coded in your function. That makes it very hard to unit test. So the first step would be to extract the URL from the function. Without knowing the rest of you code, Handler seems like a nice place to do this. Simplified:

type Handler struct {
    backend *url.URL
}

func NewHandler() (*Handler, error) {
    backend, err := url.Parse("http://test.com")
    if err != nil {
        return nil, err
    }
    return &Handler{backend}, nil
}

func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
    reverseProxy := &httputil.ReverseProxy{
        Director: func(r *http.Request) {
            r.URL.Host = h.backend.Host
            r.URL.Path = "/"
            r.URL.Scheme = h.backend.Scheme
            r.Host = url.Host
            r.Header.Set("X-Forwarded-Host", r.Header.Get("Host"))      
        },
    }
    reverseProxy.ServeHTTP(w, r)
}

Note that I have renamed Forward to ServeHTTP to simplify this example.

2. Use httptest for live handler testing

The next step is to have a basic test:

func TestHandler(t *testing.T) {
    // 1. set-up a backend server
    // 2. set-up a reverse proxy with the handler we are testing
    // 3. call the reverse-proxy
    // 4. check that the backend server received the correct header

}

Let's start by filling in the simple parts:

// set-up a backend server 
backendServer := httptest.NewServer(http.DefaultServeMux)
defer backendServer.Close()

backendURL, err := url.Parse(backendServer.URL)
if err != nil {
    t.Fatal(err)
}

// set-up the reverse proxy
handler := &Handler{backend: backendURL} // <-- here we inject our own endpoint!
reverseProxy := httptest.NewServer(handler)
defer reverseProxy.Close()

reverseProxyURL, err := url.Parse(reverseProxy.URL)
if err != nil {
    t.Fatal(err)
}

// call the reverse proxy
res, err := http.Get(reverseProxy.URL)
if err != nil {
    t.Fatal(err)
}
// todo optional: assert properties of the response
_ = res


// check that the backend server received the correct header
// this comes next...

3. Communicate results from test server to test

Now what we need is a way to communicate the received header to the main test. Since our test servers can use arbitrary handlers, let's extend the set-up of our backend server.

var (
    mu     sync.Mutex
    header string
)

// create a backend server that checks the incoming headers
backendServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    mu.Lock()
    defer mu.Unlock()
    header = r.Header.Get("X-Forwarded-Host")
    w.WriteHeader(http.StatusOK)
}))
defer backendServer.Close()

Note how I'm using a mutex, because the handler will run in a different go-routine. You could also use a channel.

At this point, we can implement our assertion:

mu.Lock()
got := header
mu.Unlock()

// check that the header has been set
want := reverseProxyURL.Host
if got != want {
    t.Errorf("GET %s gives header %s, got %s", reverseProxy.URL, want, got)
}

Note that this will still fail, but this time because your code under test is wrong :-) r.Header.Get("Host") should be replaced by r.Host.

Appendix: full example

package example

import (
    "net/http"
    "net/http/httptest"
    "net/http/httputil"
    "net/url"
    "sync"
    "testing"
)

type Handler struct {
    backend *url.URL
}

func NewHandler() (*Handler, error) {
    backend, err := url.Parse("http://test.com")
    if err != nil {
        return nil, err
    }
    return &Handler{backend}, nil
}

func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
    reverseProxy := &httputil.ReverseProxy{
        Director: func(r *http.Request) {
            r.URL.Host = h.backend.Host
            r.URL.Path = "/"
            r.URL.Scheme = h.backend.Scheme
            r.Header.Set("X-Forwarded-Host", r.Host)
            r.Host = h.backend.Host
        },
    }
    reverseProxy.ServeHTTP(w, r)
}

func TestHandler(t *testing.T) {
    var (
        mu     sync.Mutex
        header string
    )

    // create a backend server that checks the incoming headers
    backendServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        mu.Lock()
        defer mu.Unlock()
        header = r.Header.Get("X-Forwarded-Host")
        w.WriteHeader(http.StatusOK)
    }))
    defer backendServer.Close()

    backendURL, err := url.Parse(backendServer.URL)
    if err != nil {
        t.Fatal(err)
    }

    // create a server for your reverse proxy
    handler := &Handler{backend: backendURL}
    reverseProxy := httptest.NewServer(handler)
    defer reverseProxy.Close()

    reverseProxyURL, err := url.Parse(reverseProxy.URL)
    if err != nil {
        t.Fatal(err)
    }

    // make a request to the reverse proxy
    res, err := http.Get(reverseProxy.URL)
    if err != nil {
        t.Fatal(err)
    }
    // todo optional: assert properties of the response
    _ = res

    mu.Lock()
    got := header
    mu.Unlock()

    // check that the header has been set
    want := reverseProxyURL.Host
    if got != want {
        t.Errorf("GET %s gives header %s, got %s", reverseProxy.URL, want, got)
    }
}