github.com/cilium/cilium@v1.16.2/pkg/hubble/testutils/grpc.go (about)

     1  // SPDX-License-Identifier: Apache-2.0
     2  // Copyright Authors of Hubble
     3  
     4  // Copyright Authors of Cilium
     5  
     6  package testutils
     7  
     8  import (
     9  	"context"
    10  
    11  	"google.golang.org/grpc/metadata"
    12  )
    13  
    14  // FakeGRPCServerStream implements google.golang.org/grpc.ServerStream
    15  // interface for unit tests.
    16  type FakeGRPCServerStream struct {
    17  	OnSetHeader  func(metadata.MD) error
    18  	OnSendHeader func(metadata.MD) error
    19  	OnSetTrailer func(m metadata.MD)
    20  	OnContext    func() context.Context
    21  	OnSendMsg    func(m interface{}) error
    22  	OnRecvMsg    func(m interface{}) error
    23  }
    24  
    25  // SetHeader implements grpc.ServerStream.SetHeader.
    26  func (s *FakeGRPCServerStream) SetHeader(m metadata.MD) error {
    27  	if s.OnSetHeader != nil {
    28  		return s.OnSetHeader(m)
    29  	}
    30  	panic("OnSetHeader not set")
    31  }
    32  
    33  // SendHeader implements grpc.ServerStream.SendHeader.
    34  func (s *FakeGRPCServerStream) SendHeader(m metadata.MD) error {
    35  	if s.OnSendHeader != nil {
    36  		return s.OnSendHeader(m)
    37  	}
    38  	panic("OnSendHeader not set")
    39  }
    40  
    41  // SetTrailer implements grpc.ServerStream.SetTrailer.
    42  func (s *FakeGRPCServerStream) SetTrailer(m metadata.MD) {
    43  	if s.OnSetTrailer != nil {
    44  		s.OnSetTrailer(m)
    45  	}
    46  	panic("OnSetTrailer not set")
    47  }
    48  
    49  // Context implements grpc.ServerStream.Context.
    50  func (s *FakeGRPCServerStream) Context() context.Context {
    51  	if s.OnContext != nil {
    52  		return s.OnContext()
    53  	}
    54  	panic("OnContext not set")
    55  }
    56  
    57  // SendMsg implements grpc.ServerStream.SendMsg.
    58  func (s *FakeGRPCServerStream) SendMsg(m interface{}) error {
    59  	if s.OnSendMsg != nil {
    60  		return s.OnSendMsg(m)
    61  	}
    62  	panic("OnSendMsg not set")
    63  }
    64  
    65  // RecvMsg implements grpc.ServerStream.RecvMsg.
    66  func (s *FakeGRPCServerStream) RecvMsg(m interface{}) error {
    67  	if s.OnRecvMsg != nil {
    68  		return s.OnRecvMsg(m)
    69  	}
    70  	panic("OnRecvMsg not set")
    71  }
    72  
    73  // FakeGRPCClientStream implements google.golang.org/grpc.ClientStream
    74  // interface for unit tests.
    75  type FakeGRPCClientStream struct {
    76  	OnHeader    func() (metadata.MD, error)
    77  	OnTrailer   func() metadata.MD
    78  	OnCloseSend func() error
    79  	OnContext   func() context.Context
    80  	OnSendMsg   func(m interface{}) error
    81  	OnRecvMsg   func(m interface{}) error
    82  }
    83  
    84  // Header implements grpc.ClientStream.Header.
    85  func (c *FakeGRPCClientStream) Header() (metadata.MD, error) {
    86  	if c.OnHeader != nil {
    87  		return c.OnHeader()
    88  	}
    89  	panic("OnHeader not set")
    90  }
    91  
    92  // Trailer implements grpc.ClientStream.Trailer.
    93  func (c *FakeGRPCClientStream) Trailer() metadata.MD {
    94  	if c.OnTrailer != nil {
    95  		return c.OnTrailer()
    96  	}
    97  	panic("OnTrailer not set")
    98  }
    99  
   100  // CloseSend implements grpc.ClientStream.CloseSend.
   101  func (c *FakeGRPCClientStream) CloseSend() error {
   102  	if c.OnCloseSend != nil {
   103  		return c.OnCloseSend()
   104  	}
   105  	panic("OnCloseSend not set")
   106  }
   107  
   108  // Context implements grpc.ClientStream.Context.
   109  func (c *FakeGRPCClientStream) Context() context.Context {
   110  	if c.OnContext != nil {
   111  		return c.OnContext()
   112  	}
   113  	panic("OnContext not set")
   114  }
   115  
   116  // SendMsg implements grpc.ClientStream.SendMsg.
   117  func (c *FakeGRPCClientStream) SendMsg(m interface{}) error {
   118  	if c.OnSendMsg != nil {
   119  		return c.OnSendMsg(m)
   120  	}
   121  	panic("OnSendMsg not set")
   122  }
   123  
   124  // RecvMsg implements grpc.ClientStream.RecvMsg.
   125  func (c *FakeGRPCClientStream) RecvMsg(m interface{}) error {
   126  	if c.OnRecvMsg != nil {
   127  		return c.OnRecvMsg(m)
   128  	}
   129  	panic("OnRecvMsg not set")
   130  }