/*
 *
 * 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)
		}
	}
}
