trpc.group/trpc-go/trpc-go@v1.0.3/naming/loadbalance/random_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 loadbalance
    15  
    16  import (
    17  	"context"
    18  	"testing"
    19  
    20  	"trpc.group/trpc-go/trpc-go/naming/bannednodes"
    21  	"trpc.group/trpc-go/trpc-go/naming/registry"
    22  
    23  	"github.com/stretchr/testify/assert"
    24  	"github.com/stretchr/testify/require"
    25  )
    26  
    27  func TestRandomEmptyList(t *testing.T) {
    28  	b := NewRandom()
    29  	_, err := b.Select("", nil)
    30  	assert.Equal(t, err, ErrNoServerAvailable)
    31  }
    32  
    33  func TestRandomGet(t *testing.T) {
    34  	b := NewRandom()
    35  	node, err := b.Select("", []*registry.Node{testNode})
    36  	assert.Nil(t, err)
    37  	assert.Equal(t, node, testNode)
    38  }
    39  
    40  func TestRandom_SelectMandatoryBanned(t *testing.T) {
    41  	candidates := []*registry.Node{
    42  		{ServiceName: "1", Address: "1"},
    43  		{ServiceName: "2", Address: "2"},
    44  	}
    45  
    46  	ctx := context.Background()
    47  	ctx = bannednodes.NewCtx(ctx, true)
    48  
    49  	b := NewRandom()
    50  	_, err := b.Select("", candidates, WithContext(ctx))
    51  	require.Nil(t, err)
    52  
    53  	_, err = b.Select("", candidates, WithContext(ctx))
    54  	require.Nil(t, err)
    55  
    56  	_, err = b.Select("", candidates, WithContext(ctx))
    57  	require.NotNil(t, err)
    58  
    59  	nodes, mandatory, ok := bannednodes.FromCtx(ctx)
    60  	require.True(t, ok)
    61  	require.True(t, mandatory)
    62  	var n int
    63  	nodes.Range(func(*registry.Node) bool {
    64  		n++
    65  		return true
    66  	})
    67  	require.Equal(t, 2, n)
    68  }
    69  
    70  func TestRandom_SelectOptionalBanned(t *testing.T) {
    71  	candidates := []*registry.Node{
    72  		{ServiceName: "1", Address: "1"},
    73  		{ServiceName: "2", Address: "2"},
    74  	}
    75  
    76  	ctx := context.Background()
    77  	ctx = bannednodes.NewCtx(ctx, false)
    78  
    79  	b := NewRandom()
    80  	n1, err := b.Select("", candidates, WithContext(ctx))
    81  	require.Nil(t, err)
    82  
    83  	n2, err := b.Select("", candidates, WithContext(ctx))
    84  	require.Nil(t, err)
    85  	require.NotEqual(t, n1.Address, n2.Address)
    86  
    87  	_, err = b.Select("", candidates, WithContext(ctx))
    88  	require.Nil(t, err)
    89  
    90  	nodes, mandatory, ok := bannednodes.FromCtx(ctx)
    91  	require.True(t, ok)
    92  	require.False(t, mandatory)
    93  	var n int
    94  	nodes.Range(func(*registry.Node) bool {
    95  		n++
    96  		return true
    97  	})
    98  	require.Equal(t, 3, n)
    99  }