Not able to set header when calling test function with grpc.SetHeader or grpc.SendHeader

Question

I have a function which add something into the grpc header

func Dummy(ctx context.Context, request *service.Request) error {
  err := grpc.SetHeader(ctx, metadata.Pairs("key", "value"))
  return err
}

This works fine when it gets called with service.

While writing the unit test case for this, it always fails with:

rpc error: code = Internal desc = grpc: failed to fetch the stream from the context context.TODO

func TestDummy() {
 ctx := context.TODO()
 err := Dummy(ctx, &service.Request{})
}

I know I can mock grpc.SetHeader, but I don't want to do that. Is there any other way we can fix this?

Answer 1

The error happens because there isn't a server stream implementation inside context.TODO(). Instead of context.TODO() you can use grpc.NewContextWithServerTransportStream.

You still need to mock the ServerTransportStream interface though, but that's not excessively complicated. In fact, the purpose of this interface is exactly to allow testing code that uses grpc.SetHeader.

Skeleton code:

type mockServerTransportStream struct {}

func (m *mockServerTransportStream) Method() string {
    return "foo"
}

func (m *mockServerTransportStream) SetHeader(md metadata.MD) error {
    return nil
}

func (m *mockServerTransportStream) SendHeader(md metadata.MD) error {
    return nil
}

func (m *mockServerTransportStream) SetTrailer(md metadata.MD) error {
    return nil
}

func TestDummy() {
 ctx := grpc.NewContextWithServerTransportStream(&mockServerTransportStream{})
 err := Dummy(ctx, &service.Request{})
}