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