hierarchy_test.go

Documentation: google.golang.org/grpc/internal/hierarchy

     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 hierarchy
    20  
    21  import (
    22  	"testing"
    23  
    24  	"github.com/google/go-cmp/cmp"
    25  	"google.golang.org/grpc/attributes"
    26  	"google.golang.org/grpc/resolver"
    27  )
    28  
    29  func TestGet(t *testing.T) {
    30  	tests := []struct {
    31  		name string
    32  		addr resolver.Address
    33  		want []string
    34  	}{
    35  		{
    36  			name: "not set",
    37  			addr: resolver.Address{},
    38  			want: nil,
    39  		},
    40  		{
    41  			name: "set",
    42  			addr: resolver.Address{
    43  				BalancerAttributes: attributes.New(pathKey, pathValue{"a", "b"}),
    44  			},
    45  			want: []string{"a", "b"},
    46  		},
    47  	}
    48  	for _, tt := range tests {
    49  		t.Run(tt.name, func(t *testing.T) {
    50  			if got := Get(tt.addr); !cmp.Equal(got, tt.want) {
    51  				t.Errorf("Get() = %v, want %v", got, tt.want)
    52  			}
    53  		})
    54  	}
    55  }
    56  
    57  func TestSet(t *testing.T) {
    58  	tests := []struct {
    59  		name string
    60  		addr resolver.Address
    61  		path []string
    62  	}{
    63  		{
    64  			name: "before is not set",
    65  			addr: resolver.Address{},
    66  			path: []string{"a", "b"},
    67  		},
    68  		{
    69  			name: "before is set",
    70  			addr: resolver.Address{
    71  				BalancerAttributes: attributes.New(pathKey, pathValue{"before", "a", "b"}),
    72  			},
    73  			path: []string{"a", "b"},
    74  		},
    75  	}
    76  	for _, tt := range tests {
    77  		t.Run(tt.name, func(t *testing.T) {
    78  			newAddr := Set(tt.addr, tt.path)
    79  			newPath := Get(newAddr)
    80  			if !cmp.Equal(newPath, tt.path) {
    81  				t.Errorf("path after Set() = %v, want %v", newPath, tt.path)
    82  			}
    83  		})
    84  	}
    85  }
    86  
    87  func TestGroup(t *testing.T) {
    88  	tests := []struct {
    89  		name  string
    90  		addrs []resolver.Address
    91  		want  map[string][]resolver.Address
    92  	}{
    93  		{
    94  			name: "all with hierarchy",
    95  			addrs: []resolver.Address{
    96  				{Addr: "a0", BalancerAttributes: attributes.New(pathKey, pathValue{"a"})},
    97  				{Addr: "a1", BalancerAttributes: attributes.New(pathKey, pathValue{"a"})},
    98  				{Addr: "b0", BalancerAttributes: attributes.New(pathKey, pathValue{"b"})},
    99  				{Addr: "b1", BalancerAttributes: attributes.New(pathKey, pathValue{"b"})},
   100  			},
   101  			want: map[string][]resolver.Address{
   102  				"a": {
   103  					{Addr: "a0", BalancerAttributes: attributes.New(pathKey, pathValue{})},
   104  					{Addr: "a1", BalancerAttributes: attributes.New(pathKey, pathValue{})},
   105  				},
   106  				"b": {
   107  					{Addr: "b0", BalancerAttributes: attributes.New(pathKey, pathValue{})},
   108  					{Addr: "b1", BalancerAttributes: attributes.New(pathKey, pathValue{})},
   109  				},
   110  			},
   111  		},
   112  		{
   113  			// Addresses without hierarchy are ignored.
   114  			name: "without hierarchy",
   115  			addrs: []resolver.Address{
   116  				{Addr: "a0", BalancerAttributes: attributes.New(pathKey, pathValue{"a"})},
   117  				{Addr: "a1", BalancerAttributes: attributes.New(pathKey, pathValue{"a"})},
   118  				{Addr: "b0", BalancerAttributes: nil},
   119  				{Addr: "b1", BalancerAttributes: nil},
   120  			},
   121  			want: map[string][]resolver.Address{
   122  				"a": {
   123  					{Addr: "a0", BalancerAttributes: attributes.New(pathKey, pathValue{})},
   124  					{Addr: "a1", BalancerAttributes: attributes.New(pathKey, pathValue{})},
   125  				},
   126  			},
   127  		},
   128  		{
   129  			// If hierarchy is set to a wrong type (which should never happen),
   130  			// the address is ignored.
   131  			name: "wrong type",
   132  			addrs: []resolver.Address{
   133  				{Addr: "a0", BalancerAttributes: attributes.New(pathKey, pathValue{"a"})},
   134  				{Addr: "a1", BalancerAttributes: attributes.New(pathKey, pathValue{"a"})},
   135  				{Addr: "b0", BalancerAttributes: attributes.New(pathKey, "b")},
   136  				{Addr: "b1", BalancerAttributes: attributes.New(pathKey, 314)},
   137  			},
   138  			want: map[string][]resolver.Address{
   139  				"a": {
   140  					{Addr: "a0", BalancerAttributes: attributes.New(pathKey, pathValue{})},
   141  					{Addr: "a1", BalancerAttributes: attributes.New(pathKey, pathValue{})},
   142  				},
   143  			},
   144  		},
   145  	}
   146  	for _, tt := range tests {
   147  		t.Run(tt.name, func(t *testing.T) {
   148  			if got := Group(tt.addrs); !cmp.Equal(got, tt.want, cmp.AllowUnexported(attributes.Attributes{})) {
   149  				t.Errorf("Group() = %v, want %v", got, tt.want)
   150  				t.Errorf("diff: %v", cmp.Diff(got, tt.want, cmp.AllowUnexported(attributes.Attributes{})))
   151  			}
   152  		})
   153  	}
   154  }
   155  
   156  func TestGroupE2E(t *testing.T) {
   157  	hierarchy := map[string]map[string][]string{
   158  		"p0": {
   159  			"wt0": {"addr0", "addr1"},
   160  			"wt1": {"addr2", "addr3"},
   161  		},
   162  		"p1": {
   163  			"wt10": {"addr10", "addr11"},
   164  			"wt11": {"addr12", "addr13"},
   165  		},
   166  	}
   167  
   168  	var addrsWithHierarchy []resolver.Address
   169  	for p, wts := range hierarchy {
   170  		path1 := pathValue{p}
   171  		for wt, addrs := range wts {
   172  			path2 := append(pathValue(nil), path1...)
   173  			path2 = append(path2, wt)
   174  			for _, addr := range addrs {
   175  				a := resolver.Address{
   176  					Addr:               addr,
   177  					BalancerAttributes: attributes.New(pathKey, path2),
   178  				}
   179  				addrsWithHierarchy = append(addrsWithHierarchy, a)
   180  			}
   181  		}
   182  	}
   183  
   184  	gotHierarchy := make(map[string]map[string][]string)
   185  	for p1, wts := range Group(addrsWithHierarchy) {
   186  		gotHierarchy[p1] = make(map[string][]string)
   187  		for p2, addrs := range Group(wts) {
   188  			for _, addr := range addrs {
   189  				gotHierarchy[p1][p2] = append(gotHierarchy[p1][p2], addr.Addr)
   190  			}
   191  		}
   192  	}
   193  
   194  	if !cmp.Equal(gotHierarchy, hierarchy) {
   195  		t.Errorf("diff: %v", cmp.Diff(gotHierarchy, hierarchy))
   196  	}
   197  }
   198
View as plain text