github.com/rudderlabs/rudder-go-kit@v0.30.0/tcpproxy/tcpproxy_test.go (about)

     1  package tcpproxy
     2  
     3  import (
     4  	"io"
     5  	"net"
     6  	"net/http"
     7  	"net/http/httptest"
     8  	"net/url"
     9  	"strconv"
    10  	"testing"
    11  	"time"
    12  
    13  	"github.com/stretchr/testify/require"
    14  
    15  	"github.com/rudderlabs/rudder-go-kit/httputil"
    16  	"github.com/rudderlabs/rudder-go-kit/testhelper"
    17  )
    18  
    19  func TestProxy(t *testing.T) {
    20  	httpServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    21  		_, _ = w.Write([]byte("Hello, world!"))
    22  	}))
    23  	defer httpServer.Close()
    24  
    25  	httpServerParsedURL, err := url.Parse(httpServer.URL)
    26  	require.NoError(t, err)
    27  
    28  	_, httpServerPort, err := net.SplitHostPort(httpServerParsedURL.Host)
    29  	require.NoError(t, err)
    30  
    31  	proxyPort, err := testhelper.GetFreePort()
    32  	require.NoError(t, err)
    33  
    34  	proxy := &Proxy{
    35  		LocalAddr:  "localhost:" + strconv.Itoa(proxyPort),
    36  		RemoteAddr: "localhost:" + httpServerPort,
    37  	}
    38  	go proxy.Start(t)
    39  	t.Cleanup(proxy.Stop)
    40  
    41  	var (
    42  		body       []byte
    43  		statusCode int
    44  	)
    45  	require.Eventually(t, func() bool {
    46  		resp, err := http.Get("http://" + proxy.LocalAddr)
    47  		defer func() { httputil.CloseResponse(resp) }()
    48  		if err == nil {
    49  			statusCode = resp.StatusCode
    50  			body, err = io.ReadAll(resp.Body)
    51  		}
    52  		return err == nil
    53  	}, 5*time.Second, 10*time.Millisecond, "failed to connect to proxy")
    54  
    55  	require.Equal(t, http.StatusOK, statusCode)
    56  	require.Equal(t, "Hello, world!", string(body))
    57  }