如何在ReverseProxy中测试标题？

Question

我正在尝试对以下代码进行单元测试：

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

我无法弄清楚如何测试Director函数是否正在修改标头。我们如何在Go中的反向代理中测试标头？

Answer 1

1。将外部依赖项注入到被测单元中

我现在看到的最大问题是，转发给您的URL在功能中进行了硬编码。这使得很难进行单元测试。因此，第一步将是从函数中提取URL。在不了解其余代码的情况下，Handler似乎是执行此操作的好地方。简化：

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

请注意，为了简化此示例，我将Forward重命名为ServeHTTP。

2。使用httptest进行实时处理程序测试

下一步是进行基本测试：

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

}

让我们先填充简单的部分：

// 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。从测试服务器传达结果以进行测试

现在，我们需要的是一种将接收到的标头传达给主测试的方法。由于我们的测试服务器可以使用任意处理程序，因此让我们扩展后端服务器的设置。

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()

请注意我如何使用互斥锁，因为处理程序将在不同的例程中运行。您也可以使用频道。

此时，我们可以实现我们的断言：

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

请注意，这仍然会失败，但是这次是因为您的测试代码不正确：-) r.Header.Get("Host")应该替换为r.Host。

附录：完整示例

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