wrappers_test.go

Documentation: github.com/grpc-ecosystem/go-grpc-middleware

     1  // Copyright 2016 Michal Witkowski. All Rights Reserved.
     2  // See LICENSE for licensing terms.
     3  
     4  package grpc_middleware
     5  
     6  import (
     7  	"context"
     8  	"testing"
     9  
    10  	"github.com/stretchr/testify/assert"
    11  	"google.golang.org/grpc"
    12  	"google.golang.org/grpc/codes"
    13  	"google.golang.org/grpc/status"
    14  )
    15  
    16  func TestWrapServerStream(t *testing.T) {
    17  	ctx := context.WithValue(context.TODO(), "something", 1)
    18  	fake := &fakeServerStream{ctx: ctx}
    19  	wrapped := WrapServerStream(fake)
    20  	assert.NotNil(t, wrapped.Context().Value("something"), "values from fake must propagate to wrapper")
    21  	wrapped.WrappedContext = context.WithValue(wrapped.Context(), "other", 2)
    22  	assert.NotNil(t, wrapped.Context().Value("other"), "values from wrapper must be set")
    23  }
    24  
    25  type fakeServerStream struct {
    26  	grpc.ServerStream
    27  	ctx         context.Context
    28  	recvMessage interface{}
    29  	sentMessage interface{}
    30  }
    31  
    32  func (f *fakeServerStream) Context() context.Context {
    33  	return f.ctx
    34  }
    35  
    36  func (f *fakeServerStream) SendMsg(m interface{}) error {
    37  	if f.sentMessage != nil {
    38  		return status.Errorf(codes.AlreadyExists, "fakeServerStream only takes one message, sorry")
    39  	}
    40  	f.sentMessage = m
    41  	return nil
    42  }
    43  
    44  func (f *fakeServerStream) RecvMsg(m interface{}) error {
    45  	if f.recvMessage == nil {
    46  		return status.Errorf(codes.NotFound, "fakeServerStream has no message, sorry")
    47  	}
    48  	return nil
    49  }
    50  
    51  type fakeClientStream struct {
    52  	grpc.ClientStream
    53  }
    54
View as plain text