go.uber.org/yarpc@v1.72.1/internal/observability/common_test.go (about)

     1  // Copyright (c) 2022 Uber Technologies, Inc.
     2  //
     3  // Permission is hereby granted, free of charge, to any person obtaining a copy
     4  // of this software and associated documentation files (the "Software"), to deal
     5  // in the Software without restriction, including without limitation the rights
     6  // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
     7  // copies of the Software, and to permit persons to whom the Software is
     8  // furnished to do so, subject to the following conditions:
     9  //
    10  // The above copyright notice and this permission notice shall be included in
    11  // all copies or substantial portions of the Software.
    12  //
    13  // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
    14  // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
    15  // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
    16  // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
    17  // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
    18  // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
    19  // THE SOFTWARE.
    20  
    21  package observability
    22  
    23  import (
    24  	"bytes"
    25  	"context"
    26  	"io/ioutil"
    27  	"time"
    28  
    29  	"go.uber.org/yarpc/api/transport"
    30  	"go.uber.org/yarpc/yarpcerrors"
    31  )
    32  
    33  type fakeAck struct{}
    34  
    35  func (a fakeAck) String() string { return "" }
    36  
    37  type fakeHandler struct {
    38  	err                   error
    39  	applicationErr        bool
    40  	applicationErrName    string
    41  	applicationErrDetails string
    42  	applicationErrCode    *yarpcerrors.Code
    43  	applicationPanic      bool
    44  	handleStream          func(*transport.ServerStream)
    45  	responseData          []byte
    46  }
    47  
    48  func (h fakeHandler) Handle(_ context.Context, _ *transport.Request, rw transport.ResponseWriter) error {
    49  	if h.applicationPanic {
    50  		panic("application panicked")
    51  	}
    52  	if h.applicationErr {
    53  		rw.SetApplicationError()
    54  
    55  		if applicationErrorMetaSetter, ok := rw.(transport.ApplicationErrorMetaSetter); ok {
    56  			applicationErrorMetaSetter.SetApplicationErrorMeta(&transport.ApplicationErrorMeta{
    57  				Details: h.applicationErrDetails,
    58  				Name:    h.applicationErrName,
    59  				Code:    h.applicationErrCode,
    60  			})
    61  		}
    62  	}
    63  
    64  	if h.responseData != nil {
    65  		rw.Write(h.responseData)
    66  	}
    67  
    68  	return h.err
    69  }
    70  
    71  func (h fakeHandler) HandleOneway(context.Context, *transport.Request) error {
    72  	if h.applicationPanic {
    73  		panic("application panicked")
    74  	}
    75  	return h.err
    76  }
    77  
    78  func (h fakeHandler) HandleStream(stream *transport.ServerStream) error {
    79  	if h.applicationPanic {
    80  		panic("application panicked")
    81  	}
    82  	if h.handleStream != nil {
    83  		h.handleStream(stream)
    84  	}
    85  	return h.err
    86  }
    87  
    88  type fakeOutbound struct {
    89  	transport.Outbound
    90  
    91  	err                   error
    92  	applicationErr        bool
    93  	applicationErrName    string
    94  	applicationErrDetails string
    95  	applicationErrCode    *yarpcerrors.Code
    96  	applicationPanic      bool
    97  	stream                fakeStream
    98  
    99  	body []byte
   100  }
   101  
   102  func (o fakeOutbound) Call(context.Context, *transport.Request) (*transport.Response, error) {
   103  	if o.applicationPanic {
   104  		panic("application panicked")
   105  	}
   106  
   107  	return &transport.Response{
   108  		ApplicationError: o.applicationErr,
   109  		ApplicationErrorMeta: &transport.ApplicationErrorMeta{
   110  			Details: o.applicationErrDetails,
   111  			Name:    o.applicationErrName,
   112  			Code:    o.applicationErrCode,
   113  		},
   114  		Body:     ioutil.NopCloser(bytes.NewReader(o.body)),
   115  		BodySize: len(o.body)}, o.err
   116  }
   117  
   118  func (o fakeOutbound) CallOneway(context.Context, *transport.Request) (transport.Ack, error) {
   119  	if o.applicationPanic {
   120  		panic("application panicked")
   121  	}
   122  
   123  	if o.err != nil {
   124  		return nil, o.err
   125  	}
   126  	return fakeAck{}, nil
   127  }
   128  
   129  func (o fakeOutbound) CallStream(ctx context.Context, request *transport.StreamRequest) (*transport.ClientStream, error) {
   130  	if o.applicationPanic {
   131  		panic("application panicked")
   132  	}
   133  
   134  	if o.err != nil {
   135  		return nil, o.err
   136  	}
   137  
   138  	// always use passed in context and request
   139  	stream := o.stream
   140  	stream.ctx = ctx
   141  	stream.request = request
   142  
   143  	return transport.NewClientStream(&stream)
   144  }
   145  
   146  type fakeStream struct {
   147  	ctx     context.Context
   148  	request *transport.StreamRequest
   149  
   150  	receiveMsg *transport.StreamMessage
   151  
   152  	sendErr    error
   153  	receiveErr error
   154  	closeErr   error
   155  }
   156  
   157  func (s *fakeStream) Context() context.Context {
   158  	return s.ctx
   159  }
   160  
   161  func (s *fakeStream) Request() *transport.StreamRequest {
   162  	return s.request
   163  }
   164  
   165  func (s *fakeStream) SendMessage(context.Context, *transport.StreamMessage) error {
   166  	return s.sendErr
   167  }
   168  
   169  func (s *fakeStream) ReceiveMessage(context.Context) (*transport.StreamMessage, error) {
   170  	return s.receiveMsg, s.receiveErr
   171  }
   172  
   173  func (s *fakeStream) Close(context.Context) error {
   174  	return s.closeErr
   175  }
   176  
   177  func stubTime() func() {
   178  	return stubTimeWithTimeVal(time.Time{})
   179  }
   180  
   181  func stubTimeWithTimeVal(timeVal time.Time) func() {
   182  	prev := _timeNow
   183  	_timeNow = func() time.Time { return timeVal }
   184  	return func() { _timeNow = prev }
   185  }