Improve reliability of SDAM error 'before handshake completes' tests. (#1095)

Author: matthewdale

mongo/integration/sdam_error_handling_test.go

@@ -85,7 +85,23 @@ func TestSDAMErrorHandling(t *testing.T) {
 				assert.NotNil(mt, err, "expected InsertOne error, got nil")
 				assert.True(mt, mongo.IsTimeout(err), "expected timeout error, got %v", err)
 				assert.True(mt, mongo.IsNetworkError(err), "expected network error, got %v", err)
-				assert.True(mt, tpm.IsPoolCleared(), "expected pool to be cleared but was not")
+				// Assert that the pool is cleared within 2 seconds.
+				assert.Soon(mt, func(ctx context.Context) {
+					ticker := time.NewTicker(100 * time.Millisecond)
+					defer ticker.Stop()
+
+					for {
+						select {
+						case <-ticker.C:
+						case <-ctx.Done():
+							return
+						}
+
+						if tpm.IsPoolCleared() {
+							return
+						}
+					}
+				}, 2*time.Second)
 			})
 
 			mt.RunOpts("pool cleared on non-timeout network error", noClientOpts, func(mt *mtest.T) {
@@ -114,9 +130,23 @@ func TestSDAMErrorHandling(t *testing.T) {
 						// Set minPoolSize to enable the background pool maintenance goroutine.
 						SetMinPoolSize(5))
 
-					time.Sleep(200 * time.Millisecond)
-
-					assert.True(mt, tpm.IsPoolCleared(), "expected pool to be cleared but was not")
+					// Assert that the pool is cleared within 2 seconds.
+					assert.Soon(mt, func(ctx context.Context) {
+						ticker := time.NewTicker(100 * time.Millisecond)
+						defer ticker.Stop()
+
+						for {
+							select {
+							case <-ticker.C:
+							case <-ctx.Done():
+								return
+							}
+
+							if tpm.IsPoolCleared() {
+								return
+							}
+						}
+					}, 2*time.Second)
 				})
 
 				mt.Run("foreground", func(mt *mtest.T) {
@@ -143,7 +173,24 @@ func TestSDAMErrorHandling(t *testing.T) {
 					_, err := mt.Coll.InsertOne(context.Background(), bson.D{{"x", 1}})
 					assert.NotNil(mt, err, "expected InsertOne error, got nil")
 					assert.False(mt, mongo.IsTimeout(err), "expected non-timeout error, got %v", err)
-					assert.True(mt, tpm.IsPoolCleared(), "expected pool to be cleared but was not")
+
+					// Assert that the pool is cleared within 2 seconds.
+					assert.Soon(mt, func(ctx context.Context) {
+						ticker := time.NewTicker(100 * time.Millisecond)
+						defer ticker.Stop()
+
+						for {
+							select {
+							case <-ticker.C:
+							case <-ctx.Done():
+								return
+							}
+
+							if tpm.IsPoolCleared() {
+								return
+							}
+						}
+					}, 2*time.Second)
 				})
 			})
 		})