@@ -11,6 +11,19 @@ import (
1111 "github.com/stretchr/testify/require"
1212)
1313
14+ // pollUntil polls a condition until it's true or timeout
15+ func pollUntil (t * testing.T , condition func () bool , timeout time.Duration , message string ) {
16+ t .Helper ()
17+ deadline := time .Now ().Add (timeout )
18+ for time .Now ().Before (deadline ) {
19+ if condition () {
20+ return
21+ }
22+ time .Sleep (5 * time .Millisecond )
23+ }
24+ t .Fatal (message )
25+ }
26+
1427func TestNewCircuitBreaker (t * testing.T ) {
1528 tests := []struct {
1629 name string
@@ -176,16 +189,18 @@ func TestCircuitBreaker_RecoveryFlow(t *testing.T) {
176189 cb .RecordFailure ("test-nodeclass" , "us-south" , fmt .Errorf ("test error 2" ))
177190 assert .Equal (t , CircuitBreakerOpen , cb .state )
178191
179- // Wait for recovery timeout
180- time .Sleep (150 * time .Millisecond )
192+ // Wait for recovery timeout using polling
193+ pollUntil (t , func () bool {
194+ // Check if we can provision - if so, circuit breaker is in half-open
195+ err := cb .CanProvision (ctx , "test-nodeclass" , "us-south" , 0 )
196+ return err == nil
197+ }, 500 * time .Millisecond , "circuit breaker should transition to half-open" )
181198
182- // Should transition to half-open and allow limited requests
183- err := cb .CanProvision (ctx , "test-nodeclass" , "us-south" , 0 )
184- assert .NoError (t , err )
199+ // Should be in half-open state now
185200 assert .Equal (t , CircuitBreakerHalfOpen , cb .state )
186201
187202 // Second request should be blocked (exceeds HalfOpenMaxRequests)
188- err = cb .CanProvision (ctx , "test-nodeclass" , "us-south" , 0 )
203+ err : =cb .CanProvision (ctx , "test-nodeclass" , "us-south" , 0 )
189204 assert .Error (t , err )
190205 assert .IsType (t , & CircuitBreakerError {}, err )
191206
@@ -215,11 +230,11 @@ func TestCircuitBreaker_HalfOpenFailureReturnsToOpen(t *testing.T) {
215230 cb .RecordFailure ("test-nodeclass" , "us-south" , fmt .Errorf ("test error 2" ))
216231 assert .Equal (t , CircuitBreakerOpen , cb .state )
217232
218- // Wait for recovery and transition to half-open
219- time . Sleep ( 150 * time . Millisecond )
220- err := cb .CanProvision (context .Background (), "test-nodeclass" , "us-south" , 0 )
221- assert . NoError ( t , err )
222- assert . Equal ( t , CircuitBreakerHalfOpen , cb . state )
233+ // Wait for recovery and transition to half-open using polling
234+ pollUntil ( t , func () bool {
235+ err := cb .CanProvision (context .Background (), "test-nodeclass" , "us-south" , 0 )
236+ return err == nil && cb . state == CircuitBreakerHalfOpen
237+ }, 500 * time . Millisecond , "circuit breaker should transition to half-open" )
223238
224239 // Record failure in half-open state should return to open
225240 cb .RecordFailure ("test-nodeclass" , "us-south" , fmt .Errorf ("test error 3" ))
@@ -362,13 +377,16 @@ func TestCircuitBreaker_CleanOldFailures(t *testing.T) {
362377 require .NoError (t , err )
363378 assert .Equal (t , 2 , status .RecentFailures )
364379
365- // Wait for failures to age out
366- time .Sleep (150 * time .Millisecond )
380+ // Wait for failures to age out using polling
381+ pollUntil (t , func () bool {
382+ st , stErr := cb .GetState ()
383+ return stErr == nil && st .RecentFailures == 0
384+ }, 500 * time .Millisecond , "old failures should be cleaned up" )
367385
368- // Trigger cleanup by calling GetState
386+ // Verify cleanup occurred
369387 status , err = cb .GetState ()
370388 require .NoError (t , err )
371- assert .Equal (t , 0 , status .RecentFailures ) // Old failures should be cleaned up
389+ assert .Equal (t , 0 , status .RecentFailures )
372390}
373391
374392func TestCircuitBreaker_ErrorTypes (t * testing.T ) {
0 commit comments