From c474ceca153162db4acc22d5c73d0e2bf2fcbbeb Mon Sep 17 00:00:00 2001
From: Andrew Burke <31974658+atburke@users.noreply.github.com>
Date: Thu, 2 May 2024 09:02:07 -0700
Subject: [PATCH] Don't time out proxy handler (#41121)

---
 api/testhelpers/proxy.go | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/api/testhelpers/proxy.go b/api/testhelpers/proxy.go
index 3a39d561758a9..6d643fb1a3a07 100644
--- a/api/testhelpers/proxy.go
+++ b/api/testhelpers/proxy.go
@@ -21,7 +21,6 @@ import (
 	"net/http/httptest"
 	"sync"
 	"testing"
-	"time"
 
 	"github.com/gravitational/trace"
 	"github.com/stretchr/testify/require"
@@ -88,9 +87,9 @@ func (p *ProxyHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	go replicate(sconn, dconn)
 	go replicate(dconn, io.MultiReader(buf, sconn))
 
-	// Wait until done, error, or 10 second.
+	// Wait until done.
 	select {
-	case <-time.After(10 * time.Second):
+	case <-r.Context().Done():
 	case <-errc:
 	}
 }