package pkg import ( "fmt" "io" "log" "net/http" "net/http/httptest" "net/url" "sync" "testing" "time" ) func Test_Proxy_ToPrivateServer(t *testing.T) { wantBodyText := "OK" wantBody := []byte(wantBodyText) upstreamSvr := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { if r.Body != nil { defer r.Body.Close() } w.WriteHeader(http.StatusOK) w.Write(wantBody) })) defer upstreamSvr.Close() port := 8080 u, _ := url.Parse(upstreamSvr.URL) log.Println("Host", u.Host) upstreamAddr := u.Host proxy := NewProxy(upstreamAddr, 8080, "127.0.0.1", time.Second*1, &mockResolver{}) gwChan := make(chan string, 1) doneCh := make(chan bool) go proxy.Start() wg := sync.WaitGroup{} wg.Add(1) go func() { gwChan <- u.Host wg.Done() }() wg.Wait() req, err := http.NewRequest(http.MethodGet, fmt.Sprintf("http://127.0.0.1:%d", port), nil) if err != nil { t.Fatal(err) } for i := 1; i < 11; i++ { res, err := http.DefaultClient.Do(req) if err != nil { t.Logf("Try %d, gave error: %s", i, err) time.Sleep(time.Millisecond * 100) } else { resBody, _ := io.ReadAll(res.Body) if string(resBody) != string(wantBody) { t.Errorf("want %s, but got %s in body", string(wantBody), string(resBody)) } break } } go func() { doneCh <- true }() } type mockResolver struct { } func (m *mockResolver) Start() { } func (m *mockResolver) Get(upstream string, got chan<- string, timeout time.Duration) { got <- upstream }