go.uber.org/yarpc@v1.72.1/encoding/thrift/multiplex_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 thrift
    22  
    23  import (
    24  	"bytes"
    25  	"strings"
    26  	"testing"
    27  
    28  	"github.com/golang/mock/gomock"
    29  	"github.com/stretchr/testify/assert"
    30  	"github.com/stretchr/testify/require"
    31  	"go.uber.org/thriftrw/protocol/stream"
    32  	"go.uber.org/thriftrw/thrifttest"
    33  	"go.uber.org/thriftrw/thrifttest/streamtest"
    34  	"go.uber.org/thriftrw/wire"
    35  )
    36  
    37  func TestMultiplexedEncode(t *testing.T) {
    38  	mockCtrl := gomock.NewController(t)
    39  	defer mockCtrl.Finish()
    40  
    41  	tests := []struct {
    42  		service  string
    43  		giveName string
    44  		wantName string
    45  	}{
    46  		{
    47  			service:  "Foo",
    48  			giveName: "bar",
    49  			wantName: "Foo:bar",
    50  		},
    51  		{
    52  			service:  "",
    53  			giveName: "y",
    54  			wantName: ":y",
    55  		},
    56  	}
    57  
    58  	for _, tt := range tests {
    59  		testName := strings.Join([]string{tt.service, tt.giveName, tt.wantName}, "-")
    60  
    61  		t.Run("wireproto-"+testName, func(t *testing.T) {
    62  			mockProto := thrifttest.NewMockProtocol(mockCtrl)
    63  			proto := multiplexedOutboundProtocol{
    64  				Protocol: mockProto,
    65  				Service:  tt.service,
    66  			}
    67  
    68  			giveEnvelope := wire.Envelope{
    69  				Value: wire.NewValueStruct(wire.Struct{Fields: []wire.Field{}}),
    70  				Type:  wire.Call,
    71  				Name:  tt.giveName,
    72  				SeqID: 42,
    73  			}
    74  
    75  			wantEnvelope := giveEnvelope
    76  			wantEnvelope.Name = tt.wantName
    77  			mockProto.EXPECT().EncodeEnveloped(wantEnvelope, gomock.Any()).Return(nil)
    78  
    79  			assert.NoError(t, proto.EncodeEnveloped(giveEnvelope, new(bytes.Buffer)))
    80  		})
    81  
    82  		t.Run("nowireproto-"+testName, func(t *testing.T) {
    83  			mockStreamWriter := streamtest.NewMockWriter(mockCtrl)
    84  			streamWriter := multiplexedNoWireWriter{
    85  				Writer:  mockStreamWriter,
    86  				Service: tt.service,
    87  			}
    88  
    89  			giveEnvelope := stream.EnvelopeHeader{
    90  				Type:  wire.Call,
    91  				Name:  tt.giveName,
    92  				SeqID: 42,
    93  			}
    94  
    95  			wantEnvelope := giveEnvelope
    96  			wantEnvelope.Name = tt.wantName
    97  			mockStreamWriter.EXPECT().WriteEnvelopeBegin(wantEnvelope).Return(nil)
    98  
    99  			assert.NoError(t, streamWriter.WriteEnvelopeBegin(giveEnvelope))
   100  		})
   101  	}
   102  }
   103  
   104  func TestMultiplexedDecode(t *testing.T) {
   105  	mockCtrl := gomock.NewController(t)
   106  	defer mockCtrl.Finish()
   107  
   108  	tests := []struct {
   109  		service  string
   110  		giveName string
   111  		wantName string
   112  	}{
   113  		{
   114  			service:  "Foo",
   115  			giveName: "Foo:bar",
   116  			wantName: "bar",
   117  		},
   118  		{
   119  			service:  "Foo",
   120  			giveName: "Bar:baz",
   121  			wantName: "Bar:baz",
   122  		},
   123  	}
   124  
   125  	for _, tt := range tests {
   126  		testName := strings.Join([]string{tt.service, tt.giveName, tt.wantName}, "-")
   127  
   128  		t.Run("wireproto-"+testName, func(t *testing.T) {
   129  			mockProto := thrifttest.NewMockProtocol(mockCtrl)
   130  			proto := multiplexedOutboundProtocol{
   131  				Protocol: mockProto,
   132  				Service:  tt.service,
   133  			}
   134  
   135  			mockProto.EXPECT().DecodeEnveloped(gomock.Any()).Return(
   136  				wire.Envelope{
   137  					Value: wire.NewValueStruct(wire.Struct{Fields: []wire.Field{}}),
   138  					Type:  wire.Call,
   139  					Name:  tt.giveName,
   140  					SeqID: 42,
   141  				}, nil)
   142  
   143  			gotEnvelope, err := proto.DecodeEnveloped(bytes.NewReader([]byte{}))
   144  			if assert.NoError(t, err) {
   145  				assert.Equal(t, tt.wantName, gotEnvelope.Name)
   146  			}
   147  		})
   148  
   149  		t.Run("nowireproto-"+testName, func(t *testing.T) {
   150  			mockStreamReader := streamtest.NewMockReader(mockCtrl)
   151  			streamReader := multiplexedNoWireReader{
   152  				Reader:  mockStreamReader,
   153  				Service: tt.service,
   154  			}
   155  
   156  			mockStreamReader.EXPECT().ReadEnvelopeBegin().Return(
   157  				stream.EnvelopeHeader{
   158  					Type:  wire.Call,
   159  					Name:  tt.giveName,
   160  					SeqID: 42,
   161  				}, nil)
   162  
   163  			gotEnvelope, err := streamReader.ReadEnvelopeBegin()
   164  			require.NoError(t, err)
   165  			assert.Equal(t, tt.wantName, gotEnvelope.Name)
   166  		})
   167  	}
   168  }