go.uber.org/yarpc@v1.72.1/internal/grpcctx/grpcctx_test.go (about)

     1  // Copyright (c) 2022 Uber Technologies, Inc.
     2  //
     3  // Permission is hereby granted, free of charge, to any person obtaining a copy
     4  // of this software and associated documentation files (the "Software"), to deal
     5  // in the Software without restriction, including without limitation the rights
     6  // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
     7  // copies of the Software, and to permit persons to whom the Software is
     8  // furnished to do so, subject to the following conditions:
     9  //
    10  // The above copyright notice and this permission notice shall be included in
    11  // all copies or substantial portions of the Software.
    12  //
    13  // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
    14  // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
    15  // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
    16  // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
    17  // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
    18  // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
    19  // THE SOFTWARE.
    20  
    21  package grpcctx
    22  
    23  import (
    24  	"context"
    25  	"testing"
    26  
    27  	"github.com/stretchr/testify/assert"
    28  	"google.golang.org/grpc/metadata"
    29  )
    30  
    31  func TestContextWrapper(t *testing.T) {
    32  	t.Run("NewContextWrapper", func(t *testing.T) {
    33  		testContextWrapper(t, NewContextWrapper)
    34  	})
    35  	t.Run("&ContextWrapper{}", func(t *testing.T) {
    36  		testContextWrapper(t, func() *ContextWrapper { return &ContextWrapper{} })
    37  	})
    38  }
    39  
    40  func testContextWrapper(t *testing.T, create func() *ContextWrapper) {
    41  	c1 := create().WithCaller("test-caller")
    42  	checkMetadata(t, c1, metadata.MD{
    43  		"rpc-caller": []string{"test-caller"},
    44  	})
    45  	c2 := c1.WithService("test-service")
    46  	checkMetadata(t, c1, metadata.MD{
    47  		"rpc-caller": []string{"test-caller"},
    48  	})
    49  	checkMetadata(t, c2, metadata.MD{
    50  		"rpc-caller":  []string{"test-caller"},
    51  		"rpc-service": []string{"test-service"},
    52  	})
    53  	c2 = c2.WithShardKey("test-shard-key").
    54  		WithRoutingKey("test-routing-key").
    55  		WithRoutingDelegate("test-routing-delegate").
    56  		WithEncoding("test-encoding")
    57  	checkMetadata(t, c2, metadata.MD{
    58  		"rpc-caller":           []string{"test-caller"},
    59  		"rpc-service":          []string{"test-service"},
    60  		"rpc-shard-key":        []string{"test-shard-key"},
    61  		"rpc-routing-key":      []string{"test-routing-key"},
    62  		"rpc-routing-delegate": []string{"test-routing-delegate"},
    63  		"rpc-encoding":         []string{"test-encoding"},
    64  	})
    65  }
    66  
    67  func checkMetadata(t *testing.T, contextWrapper *ContextWrapper, expectedMD metadata.MD) {
    68  	md, ok := metadata.FromOutgoingContext(contextWrapper.Wrap(context.Background()))
    69  	assert.True(t, ok)
    70  	assert.Equal(t, expectedMD, md)
    71  }