@@ -520,17 +520,14 @@ func (s) TestControlChannelConnectivityStateTransitions(t *testing.T) {
520520 // Setup callback to count invocations
521521 var mu sync.Mutex
522522 var callbackCount int
523- // Buffered channel to collect callback invocations without blocking
524- callbackInvoked := make (chan struct {}, tt . wantCallbackCount + 5 )
523+ // Buffered channel large enough to never block
524+ callbackInvoked := make (chan struct {}, 100 )
525525 callback := func () {
526526 mu .Lock ()
527527 callbackCount ++
528528 mu .Unlock ()
529- // Non-blocking send - if channel is full, we still counted it
530- select {
531- case callbackInvoked <- struct {}{}:
532- default :
533- }
529+ // Send to channel - should never block with large buffer
530+ callbackInvoked <- struct {}{}
534531 }
535532
536533 // Create control channel
@@ -540,57 +537,46 @@ func (s) TestControlChannelConnectivityStateTransitions(t *testing.T) {
540537 }
541538 defer ctrlCh .close ()
542539
543- // Wait for initial READY state using state change notifications
540+ // Wait for initial connection to be established and monitoring to start
541+ // We need the monitoring goroutine to process the first READY state
542+ // before we inject test states
544543 ctx , cancel := context .WithTimeout (context .Background (), defaultTestTimeout )
545544 defer cancel ()
546545
547- readyCh := make (chan struct {})
548- go func () {
549- for {
550- state := ctrlCh .cc .GetState ()
551- if state == connectivity .Ready {
552- close (readyCh )
553- return
554- }
555- if ! ctrlCh .cc .WaitForStateChange (ctx , state ) {
556- return
557- }
546+ // Wait for the ClientConn to reach READY
547+ for {
548+ state := ctrlCh .cc .GetState ()
549+ if state == connectivity .Ready {
550+ break
551+ }
552+ if ! ctrlCh .cc .WaitForStateChange (ctx , state ) {
553+ t .Fatal ("Context cancelled while waiting for READY state" )
558554 }
559- }()
560-
561- select {
562- case <- readyCh :
563- // Initial READY state achieved
564- case <- ctx .Done ():
565- t .Fatal ("Timeout waiting for initial READY state" )
566555 }
567556
568- // Process states sequentially, waiting for callbacks when expected
569- seenTransientFailure := false
570- expectedCallbacks := 0
557+ // Give monitoring goroutine time to process the initial READY state
558+ // before injecting test states
559+ time . Sleep ( 50 * time . Millisecond )
571560
561+ // Inject all test states
572562 for _ , state := range tt .states {
573- // Inject the state
574563 ctrlCh .OnMessage (state )
564+ }
575565
576- // Track if we're in a failure state
577- if state == connectivity .TransientFailure {
578- seenTransientFailure = true
579- }
566+ // Wait for all expected callbacks with timeout
567+ callbackTimeout := time .NewTimer (defaultTestTimeout )
568+ defer callbackTimeout .Stop ()
580569
581- // If transitioning to READY after a failure, wait for callback
582- if state == connectivity .Ready && seenTransientFailure {
583- expectedCallbacks ++
584- select {
585- case <- callbackInvoked :
586- // Callback received as expected
587- seenTransientFailure = false
588- case <- time .After (defaultTestTimeout ):
589- mu .Lock ()
590- got := callbackCount
591- mu .Unlock ()
592- t .Fatalf ("Timeout waiting for callback %d/%d after TRANSIENT_FAILURE→READY (got %d callbacks so far)" , expectedCallbacks , tt .wantCallbackCount , got )
593- }
570+ receivedCallbacks := 0
571+ for receivedCallbacks < tt .wantCallbackCount {
572+ select {
573+ case <- callbackInvoked :
574+ receivedCallbacks ++
575+ case <- callbackTimeout .C :
576+ mu .Lock ()
577+ got := callbackCount
578+ mu .Unlock ()
579+ t .Fatalf ("Timeout waiting for callbacks: expected %d, received %d via channel, callback count is %d" , tt .wantCallbackCount , receivedCallbacks , got )
594580 }
595581 }
596582
0 commit comments