Repository: grpc/grpc-go
File: balancer/pickfirst/pickfirstleaf/pickfirstleaf_ext_test.go (Go)

Code before:
	cc.Connect()

	testutils.AwaitState(ctx, t, cc, connectivity.Connecting)

	// Verify that the first server is contacted and the second is not.
	if holds[0].Wait(ctx) != true {
		t.Fatalf("Timeout waiting for server %d with address %q to be contacted", 0, addrs[0])
	}
	if holds[1].IsStarted() != false {
		t.Fatalf("Server %d with address %q contacted unexpectedly", 1, addrs[1])
	}

	timerCh <- struct{}{}

	// Second connection attempt is successful.
	if holds[1].Wait(ctx) != true {
		t.Fatalf("Timeout waiting for server %d with address %q to be contacted", 1, addrs[1])
	}
	holds[1].Resume()
	testutils.AwaitState(ctx, t, cc, connectivity.Ready)
}

// TestPickFirstLeaf_HappyEyeballs_TFThenTimerFires tests the pickfirst balancer
// by causing a SubConn to fail and then jumping to the 3rd SubConn after the
// happy eyeballs timer expires.
func (s) TestPickFirstLeaf_HappyEyeballs_TFThenTimerFires(t *testing.T) {
	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
	defer cancel()

	timerMu := sync.Mutex{}
	timerCh := make(chan struct{})
	originalTimer := pfinternal.TimeAfterFunc
	pfinternal.TimeAfterFunc = func(_ time.Duration, f func()) *time.Timer {
		// Set a really long expiration to prevent it from triggering
		// automatically.
		timer := time.AfterFunc(time.Hour, f)
		timerMu.Lock()
		ch := timerCh
		timerMu.Unlock()
		go func() {
			select {
			case <-ctx.Done():
			case <-ch:
			}
			timer.Reset(0)
		}()
		return timer
	}

	defer func() {
		pfinternal.TimeAfterFunc = originalTimer

Review comment (suggestion):
Here and in the first test added in this PR, could we add an `_` after the `TF`? So, this would be `TestPickFirstLeaf_HappyEyeballs_TF_ThenTimerFires`.

Code after:
		pfinternal.TimeAfterFunc = originalTimer
	}()
	triggerTimer, timeAfter := mockTimer()
	pfinternal.TimeAfterFunc = timeAfter

	dialer := testutils.NewBlockingDialer()
	opts := []grpc.DialOption{
		grpc.WithDefaultServiceConfig(fmt.Sprintf(`{"loadBalancingConfig": [{"%s":{}}]}`, pickfirstleaf.Name)),
		grpc.WithContextDialer(dialer.DialContext),
	}
	cc, rb, bm := setupPickFirstLeaf(t, 3, opts...)
	addrs := bm.resolverAddrs()
	holds := bm.holds(dialer)
	rb.UpdateState(resolver.State{Addresses: addrs})
	cc.Connect()

	testutils.AwaitState(ctx, t, cc, connectivity.Connecting)

	// Verify that only the first server is contacted.
	if holds[0].Wait(ctx) != true {
		t.Fatalf("Timeout waiting for server %d with address %q to be contacted", 0, addrs[0])
	}
	if holds[1].IsStarted() != false {
		t.Fatalf("Server %d with address %q contacted unexpectedly", 1, addrs[1])
	}
	if holds[2].IsStarted() != false {
		t.Fatalf("Server %d with address %q contacted unexpectedly", 2, addrs[2])
	}

	// First SubConn Fails.
	holds[0].Fail(fmt.Errorf("test error"))

	// Verify that only the second server is contacted.
	if holds[1].Wait(ctx) != true {
		t.Fatalf("Timeout waiting for server %d with address %q to be contacted", 1, addrs[1])
	}
	if holds[2].IsStarted() != false {
		t.Fatalf("Server %d with address %q contacted unexpectedly", 2, addrs[2])
	}

	// The happy eyeballs timer expires, pickfirst should stop waiting for
	// server[1] to report a failure/success and request the creation of a third
	// SubConn.
	triggerTimer()
	if holds[2].Wait(ctx) != true {
		t.Fatalf("Timeout waiting for server %d with address %q to be contacted", 2, addrs[2])
	}

	// Get the connection attempt to the second server to succeed and verify
	// that the channel becomes READY.
	holds[1].Resume()