github.com/johnnyeven/libtools@v0.0.0-20191126065708-61829c1adf46/courier/transport_grpc/testify/stream_mock.go (about)

     1  package testify
     2  
     3  import (
     4  	"context"
     5  
     6  	"google.golang.org/grpc"
     7  	"google.golang.org/grpc/metadata"
     8  )
     9  
    10  func NewStreamMock(codec grpc.Codec) *StreamMock {
    11  	return &StreamMock{
    12  		codec: codec,
    13  	}
    14  }
    15  
    16  type StreamMock struct {
    17  	ctx   context.Context
    18  	codec grpc.Codec
    19  	md    metadata.MD
    20  	bytes []byte
    21  }
    22  
    23  func (s *StreamMock) Context() context.Context {
    24  	return s.ctx
    25  }
    26  
    27  func (s *StreamMock) SendMsg(m interface{}) error {
    28  	bytes, err := s.codec.Marshal(m)
    29  	if err != nil {
    30  		return err
    31  	}
    32  	s.bytes = bytes
    33  	return nil
    34  }
    35  
    36  func (s *StreamMock) RecvMsg(m interface{}) error {
    37  	if s.bytes == nil {
    38  		panic("need set bytes by SendMsg before RecvMsg")
    39  	}
    40  	err := s.codec.Unmarshal(s.bytes, m)
    41  	if err != nil {
    42  		return err
    43  	}
    44  	s.bytes = nil
    45  	return nil
    46  }
    47  
    48  func (s *StreamMock) SetHeader(md metadata.MD) error {
    49  	s.md = metadata.Join(s.md, md)
    50  
    51  	// just mock to make it can be read both in or out
    52  	s.ctx = metadata.NewIncomingContext(s.ctx, md)
    53  	s.ctx = metadata.NewOutgoingContext(s.ctx, md)
    54  	return nil
    55  }
    56  
    57  func (s *StreamMock) SendHeader(metadata.MD) error {
    58  	return nil
    59  }
    60  
    61  func (s *StreamMock) SetTrailer(metadata.MD) {
    62  }