github.com/prysmaticlabs/prysm@v1.4.4/shared/grpcutils/grpcutils_test.go (about)

     1  package grpcutils
     2  
     3  import (
     4  	"context"
     5  	"encoding/json"
     6  	"strings"
     7  	"testing"
     8  
     9  	"github.com/grpc-ecosystem/grpc-gateway/v2/runtime"
    10  	"github.com/prysmaticlabs/prysm/shared/testutil/assert"
    11  	"github.com/prysmaticlabs/prysm/shared/testutil/require"
    12  	logTest "github.com/sirupsen/logrus/hooks/test"
    13  	"google.golang.org/grpc"
    14  	"google.golang.org/grpc/metadata"
    15  )
    16  
    17  type customErrorData struct {
    18  	Message string `json:"message"`
    19  }
    20  
    21  func TestAppendHeaders(t *testing.T) {
    22  	t.Run("one_header", func(t *testing.T) {
    23  		ctx := AppendHeaders(context.Background(), []string{"first=value1"})
    24  		md, ok := metadata.FromOutgoingContext(ctx)
    25  		require.Equal(t, true, ok, "Failed to read context metadata")
    26  		require.Equal(t, 1, md.Len(), "MetadataV0 contains wrong number of values")
    27  		assert.Equal(t, "value1", md.Get("first")[0])
    28  	})
    29  
    30  	t.Run("multiple_headers", func(t *testing.T) {
    31  		ctx := AppendHeaders(context.Background(), []string{"first=value1", "second=value2"})
    32  		md, ok := metadata.FromOutgoingContext(ctx)
    33  		require.Equal(t, true, ok, "Failed to read context metadata")
    34  		require.Equal(t, 2, md.Len(), "MetadataV0 contains wrong number of values")
    35  		assert.Equal(t, "value1", md.Get("first")[0])
    36  		assert.Equal(t, "value2", md.Get("second")[0])
    37  	})
    38  
    39  	t.Run("one_empty_header", func(t *testing.T) {
    40  		ctx := AppendHeaders(context.Background(), []string{"first=value1", ""})
    41  		md, ok := metadata.FromOutgoingContext(ctx)
    42  		require.Equal(t, true, ok, "Failed to read context metadata")
    43  		require.Equal(t, 1, md.Len(), "MetadataV0 contains wrong number of values")
    44  		assert.Equal(t, "value1", md.Get("first")[0])
    45  	})
    46  
    47  	t.Run("incorrect_header", func(t *testing.T) {
    48  		logHook := logTest.NewGlobal()
    49  		ctx := AppendHeaders(context.Background(), []string{"first=value1", "second"})
    50  		md, ok := metadata.FromOutgoingContext(ctx)
    51  		require.Equal(t, true, ok, "Failed to read context metadata")
    52  		require.Equal(t, 1, md.Len(), "MetadataV0 contains wrong number of values")
    53  		assert.Equal(t, "value1", md.Get("first")[0])
    54  		assert.LogsContain(t, logHook, "Skipping second")
    55  	})
    56  
    57  	t.Run("header_value_with_equal_sign", func(t *testing.T) {
    58  		ctx := AppendHeaders(context.Background(), []string{"first=value=1"})
    59  		md, ok := metadata.FromOutgoingContext(ctx)
    60  		require.Equal(t, true, ok, "Failed to read context metadata")
    61  		require.Equal(t, 1, md.Len(), "MetadataV0 contains wrong number of values")
    62  		assert.Equal(t, "value=1", md.Get("first")[0])
    63  	})
    64  }
    65  
    66  func TestAppendCustomErrorHeader(t *testing.T) {
    67  	stream := &runtime.ServerTransportStream{}
    68  	ctx := grpc.NewContextWithServerTransportStream(context.Background(), stream)
    69  	data := &customErrorData{Message: "foo"}
    70  	require.NoError(t, AppendCustomErrorHeader(ctx, data))
    71  	// The stream used in test setup sets the metadata key in lowercase.
    72  	value, ok := stream.Header()[strings.ToLower(CustomErrorMetadataKey)]
    73  	require.Equal(t, true, ok, "Failed to retrieve custom error metadata value")
    74  	expected, err := json.Marshal(data)
    75  	require.NoError(t, err)
    76  	assert.Equal(t, string(expected), value[0])
    77  
    78  }