google.golang.org/grpc@v1.72.2/balancer/pickfirst/pickfirst_test.go (about)

     1  /*
     2   *
     3   * Copyright 2024 gRPC authors.
     4   *
     5   * Licensed under the Apache License, Version 2.0 (the "License");
     6   * you may not use this file except in compliance with the License.
     7   * You may obtain a copy of the License at
     8   *
     9   *     http://www.apache.org/licenses/LICENSE-2.0
    10   *
    11   * Unless required by applicable law or agreed to in writing, software
    12   * distributed under the License is distributed on an "AS IS" BASIS,
    13   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    14   * See the License for the specific language governing permissions and
    15   * limitations under the License.
    16   *
    17   */
    18  
    19  package pickfirst
    20  
    21  import (
    22  	"context"
    23  	"errors"
    24  	"fmt"
    25  	"testing"
    26  	"time"
    27  
    28  	"google.golang.org/grpc/balancer"
    29  	"google.golang.org/grpc/connectivity"
    30  	"google.golang.org/grpc/internal/grpctest"
    31  	"google.golang.org/grpc/internal/testutils"
    32  	"google.golang.org/grpc/resolver"
    33  )
    34  
    35  const (
    36  	// Default timeout for tests in this package.
    37  	defaultTestTimeout = 10 * time.Second
    38  	// Default short timeout, to be used when waiting for events which are not
    39  	// expected to happen.
    40  	defaultTestShortTimeout = 100 * time.Millisecond
    41  )
    42  
    43  type s struct {
    44  	grpctest.Tester
    45  }
    46  
    47  func Test(t *testing.T) {
    48  	grpctest.RunSubTests(t, s{})
    49  }
    50  
    51  // TestPickFirst_InitialResolverError sends a resolver error to the balancer
    52  // before a valid resolver update. It verifies that the clientconn state is
    53  // updated to TRANSIENT_FAILURE.
    54  func (s) TestPickFirst_InitialResolverError(t *testing.T) {
    55  	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
    56  	defer cancel()
    57  	cc := testutils.NewBalancerClientConn(t)
    58  	bal := balancer.Get(Name).Build(cc, balancer.BuildOptions{})
    59  	defer bal.Close()
    60  	bal.ResolverError(errors.New("resolution failed: test error"))
    61  
    62  	if err := cc.WaitForConnectivityState(ctx, connectivity.TransientFailure); err != nil {
    63  		t.Fatalf("cc.WaitForConnectivityState(%v) returned error: %v", connectivity.TransientFailure, err)
    64  	}
    65  
    66  	// After sending a valid update, the LB policy should report CONNECTING.
    67  	ccState := balancer.ClientConnState{
    68  		ResolverState: resolver.State{
    69  			Endpoints: []resolver.Endpoint{
    70  				{Addresses: []resolver.Address{{Addr: "1.1.1.1:1"}}},
    71  				{Addresses: []resolver.Address{{Addr: "2.2.2.2:2"}}},
    72  			},
    73  		},
    74  	}
    75  	if err := bal.UpdateClientConnState(ccState); err != nil {
    76  		t.Fatalf("UpdateClientConnState(%v) returned error: %v", ccState, err)
    77  	}
    78  
    79  	if err := cc.WaitForConnectivityState(ctx, connectivity.Connecting); err != nil {
    80  		t.Fatalf("cc.WaitForConnectivityState(%v) returned error: %v", connectivity.Connecting, err)
    81  	}
    82  }
    83  
    84  // TestPickFirst_ResolverErrorinTF sends a resolver error to the balancer
    85  // before when it's attempting to connect to a SubConn TRANSIENT_FAILURE. It
    86  // verifies that the picker is updated and the SubConn is not closed.
    87  func (s) TestPickFirst_ResolverErrorinTF(t *testing.T) {
    88  	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
    89  	defer cancel()
    90  	cc := testutils.NewBalancerClientConn(t)
    91  	bal := balancer.Get(Name).Build(cc, balancer.BuildOptions{})
    92  	defer bal.Close()
    93  
    94  	// After sending a valid update, the LB policy should report CONNECTING.
    95  	ccState := balancer.ClientConnState{
    96  		ResolverState: resolver.State{
    97  			Endpoints: []resolver.Endpoint{
    98  				{Addresses: []resolver.Address{{Addr: "1.1.1.1:1"}}},
    99  			},
   100  		},
   101  	}
   102  
   103  	if err := bal.UpdateClientConnState(ccState); err != nil {
   104  		t.Fatalf("UpdateClientConnState(%v) returned error: %v", ccState, err)
   105  	}
   106  
   107  	sc1 := <-cc.NewSubConnCh
   108  	if err := cc.WaitForConnectivityState(ctx, connectivity.Connecting); err != nil {
   109  		t.Fatalf("cc.WaitForConnectivityState(%v) returned error: %v", connectivity.Connecting, err)
   110  	}
   111  
   112  	scErr := fmt.Errorf("test error: connection refused")
   113  	sc1.UpdateState(balancer.SubConnState{
   114  		ConnectivityState: connectivity.TransientFailure,
   115  		ConnectionError:   scErr,
   116  	})
   117  
   118  	if err := cc.WaitForPickerWithErr(ctx, scErr); err != nil {
   119  		t.Fatalf("cc.WaitForPickerWithErr(%v) returned error: %v", scErr, err)
   120  	}
   121  
   122  	bal.ResolverError(errors.New("resolution failed: test error"))
   123  	if err := cc.WaitForErrPicker(ctx); err != nil {
   124  		t.Fatalf("cc.WaitForPickerWithErr() returned error: %v", err)
   125  	}
   126  
   127  	select {
   128  	case <-time.After(defaultTestShortTimeout):
   129  	case sc := <-cc.ShutdownSubConnCh:
   130  		t.Fatalf("Unexpected SubConn shutdown: %v", sc)
   131  	}
   132  }