trpc.group/trpc-go/trpc-go@v1.0.3/internal/attachment/attachment_test.go (about)

     1  //
     2  //
     3  // Tencent is pleased to support the open source community by making tRPC available.
     4  //
     5  // Copyright (C) 2023 THL A29 Limited, a Tencent company.
     6  // All rights reserved.
     7  //
     8  // If you have downloaded a copy of the tRPC source code from Tencent,
     9  // please note that tRPC source code is licensed under the  Apache 2.0 License,
    10  // A copy of the Apache 2.0 License is included in this file.
    11  //
    12  //
    13  
    14  package attachment_test
    15  
    16  import (
    17  	"bytes"
    18  	"context"
    19  	"io"
    20  	"reflect"
    21  	"testing"
    22  
    23  	"github.com/stretchr/testify/require"
    24  
    25  	trpc "trpc.group/trpc-go/trpc-go"
    26  	"trpc.group/trpc-go/trpc-go/codec"
    27  	"trpc.group/trpc-go/trpc-go/internal/attachment"
    28  	"trpc.group/trpc-go/trpc-go/server"
    29  )
    30  
    31  func TestGetClientRequestAttachment(t *testing.T) {
    32  	t.Run("nil message", func(t *testing.T) {
    33  		require.Panics(t, func() {
    34  			attachment.ClientRequestAttachment(nil)
    35  		})
    36  	})
    37  	t.Run("empty message", func(t *testing.T) {
    38  		msg := trpc.Message(context.Background())
    39  		_, ok := attachment.ClientRequestAttachment(msg)
    40  		require.False(t, ok)
    41  	})
    42  	t.Run("message contains nil attachment", func(t *testing.T) {
    43  		msg := trpc.Message(context.Background())
    44  		msg.WithCommonMeta(codec.CommonMeta{attachment.ClientAttachmentKey{}: nil})
    45  		_, ok := attachment.ClientRequestAttachment(msg)
    46  		require.False(t, ok)
    47  	})
    48  	t.Run("message contains non-empty Request attachment", func(t *testing.T) {
    49  		msg := trpc.Message(context.Background())
    50  		want := bytes.NewReader([]byte("attachment"))
    51  		msg.WithCommonMeta(codec.CommonMeta{attachment.ClientAttachmentKey{}: &attachment.Attachment{Request: want}})
    52  		got, ok := attachment.ClientRequestAttachment(msg)
    53  		require.True(t, ok)
    54  		if !reflect.DeepEqual(got, want) {
    55  			t.Errorf("ServerResponseAttachment() = %v, want %v", got, want)
    56  		}
    57  	})
    58  }
    59  
    60  func TestGetServerResponseAttachment(t *testing.T) {
    61  	t.Run("nil message", func(t *testing.T) {
    62  		require.Panics(t, func() {
    63  			attachment.ServerResponseAttachment(nil)
    64  		})
    65  	})
    66  	t.Run("empty message", func(t *testing.T) {
    67  		msg := trpc.Message(context.Background())
    68  		_, ok := attachment.ServerResponseAttachment(msg)
    69  		require.False(t, ok)
    70  	})
    71  	t.Run("message contains nil attachment", func(t *testing.T) {
    72  		msg := trpc.Message(context.Background())
    73  		msg.WithCommonMeta(codec.CommonMeta{attachment.ClientAttachmentKey{}: nil})
    74  		_, ok := attachment.ClientRequestAttachment(msg)
    75  		require.False(t, ok)
    76  	})
    77  	t.Run("message contains non-empty response attachment", func(t *testing.T) {
    78  		msg := trpc.Message(context.Background())
    79  		want := bytes.NewReader([]byte("attachment"))
    80  		msg.WithCommonMeta(codec.CommonMeta{attachment.ServerAttachmentKey{}: &attachment.Attachment{Response: want}})
    81  		got, ok := attachment.ServerResponseAttachment(msg)
    82  		require.True(t, ok)
    83  		if !reflect.DeepEqual(got, want) {
    84  			t.Errorf("ServerResponseAttachment() = %v, want %v", got, want)
    85  		}
    86  	})
    87  }
    88  
    89  func TestSetClientResponseAttachment(t *testing.T) {
    90  	msg := trpc.Message(context.Background())
    91  	var a attachment.Attachment
    92  	msg.WithCommonMeta(codec.CommonMeta{attachment.ClientAttachmentKey{}: &a})
    93  	attachment.SetClientResponseAttachment(msg, []byte("attachment"))
    94  	bts, err := io.ReadAll(a.Response)
    95  
    96  	require.Nil(t, err)
    97  	require.Equal(t, []byte("attachment"), bts)
    98  }
    99  
   100  func TestSetServerAttachment(t *testing.T) {
   101  	msg := trpc.Message(context.Background())
   102  	attachment.SetServerRequestAttachment(msg, []byte("attachment"))
   103  	bts, err := io.ReadAll(server.GetAttachment(msg).Request())
   104  
   105  	require.Nil(t, err)
   106  	require.Equal(t, []byte("attachment"), bts)
   107  }