google.golang.org/grpc@v1.62.1/balancer/base/balancer_test.go (about)

     1  /*
     2   *
     3   * Copyright 2020 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 base
    20  
    21  import (
    22  	"context"
    23  	"testing"
    24  	"time"
    25  
    26  	"google.golang.org/grpc/attributes"
    27  	"google.golang.org/grpc/balancer"
    28  	"google.golang.org/grpc/connectivity"
    29  	"google.golang.org/grpc/resolver"
    30  )
    31  
    32  type testClientConn struct {
    33  	balancer.ClientConn
    34  	newSubConn func([]resolver.Address, balancer.NewSubConnOptions) (balancer.SubConn, error)
    35  }
    36  
    37  func (c *testClientConn) NewSubConn(addrs []resolver.Address, opts balancer.NewSubConnOptions) (balancer.SubConn, error) {
    38  	return c.newSubConn(addrs, opts)
    39  }
    40  
    41  func (c *testClientConn) UpdateState(balancer.State) {}
    42  
    43  type testSubConn struct {
    44  	updateState func(balancer.SubConnState)
    45  }
    46  
    47  func (sc *testSubConn) UpdateAddresses(addresses []resolver.Address) {}
    48  
    49  func (sc *testSubConn) Connect() {}
    50  
    51  func (sc *testSubConn) Shutdown() {}
    52  
    53  func (sc *testSubConn) GetOrBuildProducer(balancer.ProducerBuilder) (balancer.Producer, func()) {
    54  	return nil, nil
    55  }
    56  
    57  // testPickBuilder creates balancer.Picker for test.
    58  type testPickBuilder struct {
    59  	validate func(info PickerBuildInfo)
    60  }
    61  
    62  func (p *testPickBuilder) Build(info PickerBuildInfo) balancer.Picker {
    63  	p.validate(info)
    64  	return nil
    65  }
    66  
    67  func TestBaseBalancerReserveAttributes(t *testing.T) {
    68  	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
    69  	defer cancel()
    70  	validated := make(chan struct{}, 1)
    71  	v := func(info PickerBuildInfo) {
    72  		defer func() { validated <- struct{}{} }()
    73  		for _, sc := range info.ReadySCs {
    74  			if sc.Address.Addr == "1.1.1.1" {
    75  				if sc.Address.Attributes == nil {
    76  					t.Errorf("in picker.validate, got address %+v with nil attributes, want not nil", sc.Address)
    77  				}
    78  				foo, ok := sc.Address.Attributes.Value("foo").(string)
    79  				if !ok || foo != "2233niang" {
    80  					t.Errorf("in picker.validate, got address[1.1.1.1] with invalid attributes value %v, want 2233niang", sc.Address.Attributes.Value("foo"))
    81  				}
    82  			} else if sc.Address.Addr == "2.2.2.2" {
    83  				if sc.Address.Attributes != nil {
    84  					t.Error("in b.subConns, got address[2.2.2.2] with not nil attributes, want nil")
    85  				}
    86  			}
    87  		}
    88  	}
    89  	pickBuilder := &testPickBuilder{validate: v}
    90  	b := (&baseBuilder{pickerBuilder: pickBuilder}).Build(&testClientConn{
    91  		newSubConn: func(addrs []resolver.Address, opts balancer.NewSubConnOptions) (balancer.SubConn, error) {
    92  			return &testSubConn{updateState: opts.StateListener}, nil
    93  		},
    94  	}, balancer.BuildOptions{}).(*baseBalancer)
    95  
    96  	b.UpdateClientConnState(balancer.ClientConnState{
    97  		ResolverState: resolver.State{
    98  			Addresses: []resolver.Address{
    99  				{Addr: "1.1.1.1", Attributes: attributes.New("foo", "2233niang")},
   100  				{Addr: "2.2.2.2", Attributes: nil},
   101  			},
   102  		},
   103  	})
   104  	select {
   105  	case <-validated:
   106  	case <-ctx.Done():
   107  		t.Fatalf("timed out waiting for UpdateClientConnState to call picker.Build")
   108  	}
   109  
   110  	for sc := range b.scStates {
   111  		sc.(*testSubConn).updateState(balancer.SubConnState{ConnectivityState: connectivity.Ready, ConnectionError: nil})
   112  		select {
   113  		case <-validated:
   114  		case <-ctx.Done():
   115  			t.Fatalf("timed out waiting for UpdateClientConnState to call picker.Build")
   116  		}
   117  	}
   118  }