github.com/hxx258456/ccgo@v0.0.5-0.20230213014102-48b35f46f66f/go-grpc-middleware/wrappers_test.go (about)

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