From 17387dadc63e00f6169b842d4686fc86d733c6df Mon Sep 17 00:00:00 2001
From: MoofMonkey <11695747+MoofMonkey@users.noreply.github.com>
Date: Tue, 4 Oct 2022 06:19:37 +0300
Subject: [PATCH] Fix TestConnections pipedRequestBody flakiness

---
 proxy/proxy_test.go | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/proxy/proxy_test.go b/proxy/proxy_test.go
index 384298c3..c3a81bf1 100644
--- a/proxy/proxy_test.go
+++ b/proxy/proxy_test.go
@@ -683,9 +683,12 @@ func TestConnections(t *testing.T) {
 			req.Header = test.args.requestHeaders
 			respWriter := test.args.eyeballResponseWriter
 
+			var wg sync.WaitGroup
 			if pipedReqBody, ok := test.args.eyeballRequestBody.(*pipedRequestBody); ok {
 				respWriter = newTCPRespWriter(pipedReqBody.pipedConn)
+				wg.Add(1)
 				go func() {
+					defer wg.Done()
 					resp := pipedReqBody.roundtrip(test.args.ingressServiceScheme + ln.Addr().String())
 					replayer.Write(resp)
 				}()
@@ -698,6 +701,8 @@ func TestConnections(t *testing.T) {
 				err = proxy.ProxyHTTP(respWriter, tracing.NewTracedHTTPRequest(req, &log), test.args.connectionType == connection.TypeWebsocket)
 			}
 
+			wg.Wait()
+
 			cancel()
 			assert.Equal(t, test.want.err, err != nil)
 			assert.Equal(t, test.want.message, replayer.Bytes())