Add context to drainWarmup() in tests

Signed-off-by: Sean Chittenden <sean.chittenden@crowdstrike.com>

Author: sean-

opensearchtransport/standby_rotation_integration_test.go

@@ -142,25 +142,43 @@ func warmupSelections(rounds, skip int) int {
 
 // drainWarmup pumps enough requests through the transport to complete warmup
 // on all connections, allowing deferred cap enforcement to fire.
-func drainWarmup(transport *opensearchtransport.Client) {
+//
+// The context bounds total execution time. When the pool has 0 live connections,
+// every Perform falls through to seed-fallback (a full TLS round-trip). Without
+// a context check, the fixed iteration count can consume the entire retry budget
+// of the caller (e.g., discoverWithStandby's 30s timeout).
+//
+// Consecutive errors (maxConsecErrors) trigger an early exit: the pool is empty
+// or unhealthy and further requests would just burn time on seed fallback.
+func drainWarmup(ctx context.Context, transport *opensearchtransport.Client) {
 	const (
 		nodeCount          = 3
 		warmupRounds       = 4
 		warmupSkipMultiple = 2
+		maxConsecErrors    = 5
 	)
 	warmupSkip := warmupRounds * warmupSkipMultiple
 	selectionsPerConn := warmupSelections(warmupRounds, warmupSkip)
 	drainRequests := selectionsPerConn * nodeCount * 3 / 2
 
+	consecErrors := 0
 	for range drainRequests {
-		req, err := http.NewRequest(http.MethodGet, "/", nil)
+		if ctx.Err() != nil {
+			return
+		}
+		req, err := http.NewRequestWithContext(ctx, http.MethodGet, "/", nil)
 		if err != nil {
 			break
 		}
 		resp, err := transport.Perform(req)
 		if err != nil {
+			consecErrors++
+			if consecErrors >= maxConsecErrors {
+				return
+			}
 			continue
 		}
+		consecErrors = 0
 		resp.Body.Close()
 	}
 }
@@ -206,7 +224,7 @@ func discoverWithStandby(t *testing.T, transport *opensearchtransport.Client) op
 
 		// Pump requests to drain warmup. Cap enforcement is deferred until
 		// connections finish warming (via Next()'s deferred path).
-		drainWarmup(transport)
+		drainWarmup(ctx, transport)
 
 		m, err = transport.Metrics()
 		require.NoError(t, err)
@@ -371,7 +389,7 @@ func TestStandbyRotation(t *testing.T) {
 				time.Sleep(discoveryPause)
 				continue
 			}
-			drainWarmup(transport)
+			drainWarmup(t.Context(), transport)
 
 			promotions := obs.promotionCount()
 			demotions := obs.demotionCount()
@@ -474,7 +492,7 @@ func TestStandbyRotation(t *testing.T) {
 					time.Sleep(discoveryPause)
 					continue
 				}
-				drainWarmup(transport)
+				drainWarmup(t.Context(), transport)
 
 				promotions := obs.promotionCount()
 				t.Logf("Cycle %d attempt %d: promotions=%d (prev=%d)",