grpc-go/balancer/base/balancer_test.go

71 lines
2.0 KiB
Go

/*
*
* Copyright 2020 gRPC authors.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*
*/
package base
import (
"testing"
"google.golang.org/grpc/attributes"
"google.golang.org/grpc/balancer"
"google.golang.org/grpc/resolver"
)
type testClientConn struct {
balancer.ClientConn
newSubConn func([]resolver.Address, balancer.NewSubConnOptions) (balancer.SubConn, error)
}
func (c *testClientConn) NewSubConn(addrs []resolver.Address, opts balancer.NewSubConnOptions) (balancer.SubConn, error) {
return c.newSubConn(addrs, opts)
}
type testSubConn struct{}
func (sc *testSubConn) UpdateAddresses(addresses []resolver.Address) {}
func (sc *testSubConn) Connect() {}
func TestBaseBalancerStripAttributes(t *testing.T) {
b := (&baseBuilder{}).Build(&testClientConn{
newSubConn: func(addrs []resolver.Address, _ balancer.NewSubConnOptions) (balancer.SubConn, error) {
for _, addr := range addrs {
if addr.Attributes == nil {
t.Errorf("in NewSubConn, got address %+v with nil attributes, want not nil", addr)
}
}
return &testSubConn{}, nil
},
}, balancer.BuildOptions{}).(*baseBalancer)
b.UpdateClientConnState(balancer.ClientConnState{
ResolverState: resolver.State{
Addresses: []resolver.Address{
{Addr: "1.1.1.1", Attributes: &attributes.Attributes{}},
{Addr: "2.2.2.2", Attributes: &attributes.Attributes{}},
},
},
})
for addr := range b.subConns {
if addr.Attributes != nil {
t.Errorf("in b.subConns, got address %+v with nil attributes, want not nil", addr)
}
}
}