trpc.group/trpc-go/trpc-go@v1.0.3/naming/selector/ip_selector_test.go (about)

     1  //
     2  //
     3  // Tencent is pleased to support the open source community by making tRPC available.
     4  //
     5  // Copyright (C) 2023 THL A29 Limited, a Tencent company.
     6  // All rights reserved.
     7  //
     8  // If you have downloaded a copy of the tRPC source code from Tencent,
     9  // please note that tRPC source code is licensed under the  Apache 2.0 License,
    10  // A copy of the Apache 2.0 License is included in this file.
    11  //
    12  //
    13  
    14  package selector
    15  
    16  import (
    17  	"context"
    18  	"testing"
    19  
    20  	"github.com/stretchr/testify/assert"
    21  	"github.com/stretchr/testify/require"
    22  	"github.com/stretchr/testify/suite"
    23  
    24  	"trpc.group/trpc-go/trpc-go/naming/bannednodes"
    25  	"trpc.group/trpc-go/trpc-go/naming/registry"
    26  )
    27  
    28  // go test -v -coverprofile=cover.out
    29  // go tool cover -func=cover.out
    30  
    31  // IPSelectorTestSuite 测试 suite
    32  type IPSelectorTestSuite struct {
    33  	suite.Suite
    34  	selector Selector
    35  }
    36  
    37  // SetupSuite 初始化测试
    38  func (suite *IPSelectorTestSuite) SetupSuite() {
    39  }
    40  
    41  // 初始化测试
    42  func (suite *IPSelectorTestSuite) SetupTest() {
    43  	suite.selector = Get("ip")
    44  }
    45  
    46  func (suite *IPSelectorTestSuite) TestIPSelectSingleIP() {
    47  	serviceName := "trpc.service.ip.1:1234"
    48  	node, err := suite.selector.Select(serviceName)
    49  	suite.T().Logf("Select return node:{%+v}, err:{%+v}", node, err)
    50  
    51  	suite.NoError(err)
    52  	suite.Equal(node.ServiceName, "trpc.service.ip.1:1234")
    53  	suite.Equal(node.Address, "trpc.service.ip.1:1234")
    54  }
    55  
    56  func (suite *IPSelectorTestSuite) TestIPSelectMultiIP() {
    57  	serviceName := "trpc.service.ip.1:1234,trpc.service.ip.2:1234"
    58  
    59  	node, err := suite.selector.Select(serviceName)
    60  	suite.T().Logf("Select return node:{%+v}, err:{%+v}", node, err)
    61  	suite.NoError(err)
    62  	suite.Equal(node.ServiceName, serviceName)
    63  
    64  	node, err = suite.selector.Select(serviceName)
    65  	suite.T().Logf("Select return node:{%+v}, err:{%+v}", node, err)
    66  	suite.NoError(err)
    67  	suite.Equal(node.ServiceName, serviceName)
    68  }
    69  
    70  func (suite *IPSelectorTestSuite) TestIPSelectEmpty() {
    71  	serviceName := ""
    72  
    73  	node, err := suite.selector.Select(serviceName)
    74  	suite.T().Logf("Select return node:{%+v}, err:{%+v}", node, err)
    75  	suite.Error(err)
    76  	suite.Nil(node, "serviceName is empty")
    77  }
    78  
    79  func TestIPSelector(t *testing.T) {
    80  	suite.Run(t, new(IPSelectorTestSuite))
    81  }
    82  
    83  func TestIPSelectorSelect(t *testing.T) {
    84  	s := Get("ip")
    85  	n, err := s.Select("trpc.service.ip.1:8888")
    86  	assert.Nil(t, err)
    87  	assert.Equal(t, n.Address, "trpc.service.ip.1:8888")
    88  }
    89  
    90  func TestIPSelectorReport(t *testing.T) {
    91  	s := Get("ip")
    92  	n, err := s.Select("trpc.service.ip.1:8888")
    93  	assert.Nil(t, err)
    94  	assert.Equal(t, n.Address, "trpc.service.ip.1:8888")
    95  	assert.Nil(t, s.Report(n, 0, nil))
    96  }
    97  
    98  func TestIPSelectorSelectMandatoryBanned(t *testing.T) {
    99  	candidates := "127.0.0.1:8000,127.0.0.1:8001"
   100  
   101  	ctx := bannednodes.NewCtx(context.Background(), true)
   102  
   103  	s := NewIPSelector()
   104  
   105  	_, err := s.Select(candidates, WithContext(ctx))
   106  	require.Nil(t, err)
   107  
   108  	_, err = s.Select(candidates, WithContext(ctx))
   109  	require.Nil(t, err)
   110  
   111  	_, err = s.Select(candidates, WithContext(ctx))
   112  	require.NotNil(t, err)
   113  
   114  	nodes, mandatory, ok := bannednodes.FromCtx(ctx)
   115  	require.True(t, ok)
   116  	require.True(t, mandatory)
   117  	var n int
   118  	nodes.Range(func(*registry.Node) bool {
   119  		n++
   120  		return true
   121  	})
   122  	require.Equal(t, 2, n)
   123  }
   124  
   125  func TestIPSelectorSelectOptionalBanned(t *testing.T) {
   126  	candidates := "127.0.0.1:8000,127.0.0.1:8001"
   127  
   128  	ctx := bannednodes.NewCtx(context.Background(), false)
   129  
   130  	s := NewIPSelector()
   131  
   132  	n1, err := s.Select(candidates, WithContext(ctx))
   133  	require.Nil(t, err)
   134  
   135  	n2, err := s.Select(candidates, WithContext(ctx))
   136  	require.Nil(t, err)
   137  
   138  	require.NotEqual(t, n1.Address, n2.Address)
   139  
   140  	_, err = s.Select(candidates, WithContext(ctx))
   141  	require.Nil(t, err)
   142  
   143  	nodes, mandatory, ok := bannednodes.FromCtx(ctx)
   144  	require.True(t, ok)
   145  	require.False(t, mandatory)
   146  	var n int
   147  	nodes.Range(func(*registry.Node) bool {
   148  		n++
   149  		return true
   150  	})
   151  	require.Equal(t, 3, n)
   152  }
   153  
   154  // BenchmarkIPSelectorSelectOneService benchmark Select 性能
   155  func BenchmarkIPSelectorSelectOneService(b *testing.B) {
   156  	s := Get("ip")
   157  	for i := 0; i < b.N; i++ {
   158  		s.Select("trpc.service.ip.1:8888")
   159  	}
   160  }
   161  
   162  // BenchmarkIPSelectorSelectMultiService 测试 Select 性能
   163  func BenchmarkIPSelectorSelectMultiService(b *testing.B) {
   164  	s := Get("ip")
   165  	for i := 0; i < b.N; i++ {
   166  		s.Select("trpc.service.ip.1:8888,trpc.service.ip.1:8886,trpc.service.ip.1:8887")
   167  	}
   168  }