From 65d0dde33c2966c44cdb6eb7c1e4db2e89f6f7e5 Mon Sep 17 00:00:00 2001
From: Anton Miniailo <anton@goteleport.com>
Date: Tue, 30 Jul 2024 12:59:03 -0400
Subject: [PATCH] [v16] Fix flaky test for reissuing certificate for local kube
 proxy. (#44827)

* Fix flaky test again.

* Increase timeout to 5 seconds.
---
 lib/srv/alpnproxy/local_proxy_test.go | 18 ++++++++++++------
 1 file changed, 12 insertions(+), 6 deletions(-)

diff --git a/lib/srv/alpnproxy/local_proxy_test.go b/lib/srv/alpnproxy/local_proxy_test.go
index 1a7321457d112..f7940ef22c069 100644
--- a/lib/srv/alpnproxy/local_proxy_test.go
+++ b/lib/srv/alpnproxy/local_proxy_test.go
@@ -533,12 +533,18 @@ func TestKubeMiddleware(t *testing.T) {
 		err := km.CheckAndSetDefaults()
 		require.NoError(t, err)
 
-		rw := responsewriters.NewMemoryResponseWriter()
-		// HandleRequest will reissue certificate if needed.
-		km.HandleRequest(rw, req)
-
-		// request timed out.
-		require.Equal(t, http.StatusInternalServerError, rw.Status())
+		var rw *responsewriters.MemoryResponseWriter
+		// We use `require.Eventually` to avoid a very rare test flakiness case when reissue goroutine manages to
+		// successfully finish before the parent goroutine has a chance to check the context (and see that it's expired).
+		require.Eventually(t, func() bool {
+			rw = responsewriters.NewMemoryResponseWriter()
+			// HandleRequest will reissue certificate if needed.
+			km.HandleRequest(rw, req)
+
+			// request timed out.
+			return rw.Status() == http.StatusInternalServerError
+
+		}, 5*time.Second, 100*time.Millisecond)
 		require.Contains(t, rw.Buffer().String(), "context canceled")
 
 		// just let the reissuing goroutine some time to replace certs.