From 7b8f6272bd14b5611c67afa2e624a2ccd601c2d1 Mon Sep 17 00:00:00 2001
From: joerger <bjoerger@goteleport.com>
Date: Mon, 18 Dec 2023 11:52:03 -0800
Subject: [PATCH 1/3] Add retry to event check in
 TestTSHConfigConnectWithOpenSSHClient.

---
 tool/tsh/common/proxy_test.go | 15 ++++++++-------
 1 file changed, 8 insertions(+), 7 deletions(-)

diff --git a/tool/tsh/common/proxy_test.go b/tool/tsh/common/proxy_test.go
index 93585f7e06380..10ef2effda777 100644
--- a/tool/tsh/common/proxy_test.go
+++ b/tool/tsh/common/proxy_test.go
@@ -1100,14 +1100,15 @@ func mustSearchEvents(t *testing.T, auth *auth.Server) []apievents.AuditEvent {
 }
 
 func mustFindFailedNodeLoginAttempt(t *testing.T, s *suite, nodeLogin string) {
-	av := mustSearchEvents(t, s.root.GetAuthServer())
-	for _, e := range av {
-		if e.GetCode() == events.AuthAttemptFailureCode {
-			require.Equal(t, e.(*apievents.AuthAttempt).Login, nodeLogin)
-			return
+	require.Eventually(t, func() bool {
+		es := mustSearchEvents(t, s.root.GetAuthServer())
+		for _, e := range es {
+			if e.GetCode() == events.AuthAttemptFailureCode && e.(*apievents.AuthAttempt).Login == nodeLogin {
+				return true
+			}
 		}
-	}
-	t.Errorf("failed to find AuthAttemptFailureCode event (0/%d events matched)", len(av))
+		return false
+	}, 5*time.Second, 500*time.Millisecond, "failed to find AuthAttemptFailureCode event")
 }
 
 func TestFormatCommand(t *testing.T) {

From 5f95d87c4134eaf41c1be5630b40a6df4dfc38f0 Mon Sep 17 00:00:00 2001
From: joerger <bjoerger@goteleport.com>
Date: Tue, 19 Dec 2023 12:15:46 -0800
Subject: [PATCH 2/3] Resolve comments.

---
 tool/tsh/common/proxy_test.go | 31 +++++++++++++------------------
 1 file changed, 13 insertions(+), 18 deletions(-)

diff --git a/tool/tsh/common/proxy_test.go b/tool/tsh/common/proxy_test.go
index 10ef2effda777..ee1673ae094fa 100644
--- a/tool/tsh/common/proxy_test.go
+++ b/tool/tsh/common/proxy_test.go
@@ -1086,29 +1086,24 @@ func mustFailToRunOpenSSHCommand(t *testing.T, configFile string, sshConnString
 	require.Error(t, err)
 }
 
-func mustSearchEvents(t *testing.T, auth *auth.Server) []apievents.AuditEvent {
-	now := time.Now()
-	ctx := context.Background()
-	events, _, err := auth.SearchEvents(ctx, events.SearchEventsRequest{
-		From:  now.Add(-time.Hour),
-		To:    now.Add(time.Hour),
-		Order: types.EventOrderDescending,
-	})
-
-	require.NoError(t, err)
-	return events
-}
-
 func mustFindFailedNodeLoginAttempt(t *testing.T, s *suite, nodeLogin string) {
-	require.Eventually(t, func() bool {
-		es := mustSearchEvents(t, s.root.GetAuthServer())
+	require.EventuallyWithT(t, func(t *assert.CollectT) {
+		now := time.Now()
+		ctx := context.Background()
+		es, _, err := s.root.GetAuthServer().SearchEvents(ctx, events.SearchEventsRequest{
+			From:  now.Add(-time.Hour),
+			To:    now.Add(time.Hour),
+			Order: types.EventOrderDescending,
+		})
+		assert.NoError(t, err)
+
 		for _, e := range es {
 			if e.GetCode() == events.AuthAttemptFailureCode && e.(*apievents.AuthAttempt).Login == nodeLogin {
-				return true
+				return
 			}
 		}
-		return false
-	}, 5*time.Second, 500*time.Millisecond, "failed to find AuthAttemptFailureCode event")
+		t.Errorf("Failed to find an AuthAttemptFailureCode event")
+	}, 5*time.Second, 500*time.Millisecond)
 }
 
 func TestFormatCommand(t *testing.T) {

From e7e8fe8fb6d32c1f5dea36d42dc363377a78c341 Mon Sep 17 00:00:00 2001
From: joerger <bjoerger@goteleport.com>
Date: Wed, 20 Dec 2023 19:22:15 -0800
Subject: [PATCH 3/3] Restore nodelogin assertion.

---
 tool/tsh/common/proxy_test.go | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/tool/tsh/common/proxy_test.go b/tool/tsh/common/proxy_test.go
index ee1673ae094fa..ce68ff8e6c2df 100644
--- a/tool/tsh/common/proxy_test.go
+++ b/tool/tsh/common/proxy_test.go
@@ -1098,11 +1098,12 @@ func mustFindFailedNodeLoginAttempt(t *testing.T, s *suite, nodeLogin string) {
 		assert.NoError(t, err)
 
 		for _, e := range es {
-			if e.GetCode() == events.AuthAttemptFailureCode && e.(*apievents.AuthAttempt).Login == nodeLogin {
+			if e.GetCode() == events.AuthAttemptFailureCode {
+				assert.Equal(t, e.(*apievents.AuthAttempt).Login, nodeLogin)
 				return
 			}
 		}
-		t.Errorf("Failed to find an AuthAttemptFailureCode event")
+		t.Errorf("failed to find AuthAttemptFailureCode event (0/%d events matched)", len(es))
 	}, 5*time.Second, 500*time.Millisecond)
 }