gitee.com/ks-custle/core-gm@v0.0.0-20230922171213-b83bdd97b62c/grpc/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  	"testing"
    23  
    24  	"gitee.com/ks-custle/core-gm/grpc/attributes"
    25  	"gitee.com/ks-custle/core-gm/grpc/balancer"
    26  	"gitee.com/ks-custle/core-gm/grpc/connectivity"
    27  	"gitee.com/ks-custle/core-gm/grpc/resolver"
    28  )
    29  
    30  type testClientConn struct {
    31  	balancer.ClientConn
    32  	newSubConn func([]resolver.Address, balancer.NewSubConnOptions) (balancer.SubConn, error)
    33  }
    34  
    35  func (c *testClientConn) NewSubConn(addrs []resolver.Address, opts balancer.NewSubConnOptions) (balancer.SubConn, error) {
    36  	return c.newSubConn(addrs, opts)
    37  }
    38  
    39  func (c *testClientConn) UpdateState(balancer.State) {}
    40  
    41  type testSubConn struct{}
    42  
    43  func (sc *testSubConn) UpdateAddresses(addresses []resolver.Address) {}
    44  
    45  func (sc *testSubConn) Connect() {}
    46  
    47  // testPickBuilder creates balancer.Picker for test.
    48  type testPickBuilder struct {
    49  	validate func(info PickerBuildInfo)
    50  }
    51  
    52  func (p *testPickBuilder) Build(info PickerBuildInfo) balancer.Picker {
    53  	p.validate(info)
    54  	return nil
    55  }
    56  
    57  func TestBaseBalancerReserveAttributes(t *testing.T) {
    58  	var v = func(info PickerBuildInfo) {
    59  		for _, sc := range info.ReadySCs {
    60  			if sc.Address.Addr == "1.1.1.1" {
    61  				if sc.Address.Attributes == nil {
    62  					t.Errorf("in picker.validate, got address %+v with nil attributes, want not nil", sc.Address)
    63  				}
    64  				foo, ok := sc.Address.Attributes.Value("foo").(string)
    65  				if !ok || foo != "2233niang" {
    66  					t.Errorf("in picker.validate, got address[1.1.1.1] with invalid attributes value %v, want 2233niang", sc.Address.Attributes.Value("foo"))
    67  				}
    68  			} else if sc.Address.Addr == "2.2.2.2" {
    69  				if sc.Address.Attributes != nil {
    70  					t.Error("in b.subConns, got address[2.2.2.2] with not nil attributes, want nil")
    71  				}
    72  			}
    73  		}
    74  	}
    75  	pickBuilder := &testPickBuilder{validate: v}
    76  	b := (&baseBuilder{pickerBuilder: pickBuilder}).Build(&testClientConn{
    77  		newSubConn: func(addrs []resolver.Address, _ balancer.NewSubConnOptions) (balancer.SubConn, error) {
    78  			return &testSubConn{}, nil
    79  		},
    80  	}, balancer.BuildOptions{}).(*baseBalancer)
    81  
    82  	b.UpdateClientConnState(balancer.ClientConnState{
    83  		ResolverState: resolver.State{
    84  			Addresses: []resolver.Address{
    85  				{Addr: "1.1.1.1", Attributes: attributes.New("foo", "2233niang")},
    86  				{Addr: "2.2.2.2", Attributes: nil},
    87  			},
    88  		},
    89  	})
    90  
    91  	for sc := range b.scStates {
    92  		b.UpdateSubConnState(sc, balancer.SubConnState{ConnectivityState: connectivity.Ready, ConnectionError: nil})
    93  	}
    94  }