trpc.group/trpc-go/trpc-go@v1.0.3/naming/bannednodes/banned_nodes_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 bannednodes_test
    15  
    16  import (
    17  	"context"
    18  	"testing"
    19  
    20  	"github.com/stretchr/testify/require"
    21  
    22  	. "trpc.group/trpc-go/trpc-go/naming/bannednodes"
    23  	"trpc.group/trpc-go/trpc-go/naming/registry"
    24  )
    25  
    26  func TestNewFromCtx(t *testing.T) {
    27  	ctx := NewCtx(context.Background(), true)
    28  	_, mandatory, ok := FromCtx(ctx)
    29  	require.True(t, ok)
    30  	require.True(t, mandatory)
    31  
    32  	ctx = NewCtx(context.Background(), false)
    33  	_, mandatory, ok = FromCtx(ctx)
    34  	require.True(t, ok)
    35  	require.False(t, mandatory)
    36  }
    37  
    38  func TestBannedNodes(t *testing.T) {
    39  	const addr = "127.0.0.1:8000"
    40  
    41  	ctx := NewCtx(context.Background(), true)
    42  
    43  	nodes, _, ok := FromCtx(ctx)
    44  	require.True(t, ok)
    45  
    46  	Add(ctx, &registry.Node{Address: addr})
    47  
    48  	// the nodes returned by FromCtx is immutable, and the Add after FromCtx has no effect on nodes.
    49  	num := 0
    50  	nodes.Range(func(n *registry.Node) bool {
    51  		num++
    52  		return true
    53  	})
    54  	require.Equal(t, 0, num)
    55  
    56  	nodes, _, ok = FromCtx(ctx)
    57  	require.True(t, ok)
    58  
    59  	notfound := nodes.Range(
    60  		func(n *registry.Node) bool {
    61  			return n.Address != addr
    62  		},
    63  	)
    64  	require.False(t, notfound)
    65  }