github.com/nyan233/littlerpc@v0.4.6-0.20230316182519-0c8d5c48abaf/core/common/msgwriter/writer_test.go (about)

     1  package msgwriter
     2  
     3  import (
     4  	"github.com/nyan233/littlerpc/core/common/errorhandler"
     5  	"github.com/nyan233/littlerpc/core/common/transport"
     6  	"github.com/nyan233/littlerpc/core/container"
     7  	"github.com/nyan233/littlerpc/core/middle/packer"
     8  	"github.com/nyan233/littlerpc/core/protocol/message"
     9  	messageGen "github.com/nyan233/littlerpc/core/protocol/message/gen"
    10  	"github.com/nyan233/littlerpc/core/protocol/message/mux"
    11  	"github.com/stretchr/testify/assert"
    12  	"sync"
    13  	"syscall"
    14  	"testing"
    15  )
    16  
    17  type NilConn struct {
    18  	writeFailed bool
    19  	transport.NilConn
    20  }
    21  
    22  func (n2 *NilConn) Write(b []byte) (n int, err error) {
    23  	if n2.writeFailed {
    24  		return -1, syscall.EINPROGRESS
    25  	}
    26  	return len(b), nil
    27  }
    28  
    29  func TestLRPCWriter(t *testing.T) {
    30  	t.Run("TestLRPCNoMuxWriter", func(t *testing.T) {
    31  		testWriter(t, NewLRPCNoMux())
    32  	})
    33  	t.Run("TestLRPCMuxWriter", func(t *testing.T) {
    34  		testWriter(t, NewLRPCMux())
    35  	})
    36  	t.Run("TestJsonRPC2Writer", func(t *testing.T) {
    37  		testWriter(t, NewJsonRPC2())
    38  	})
    39  }
    40  
    41  func testWriter(t *testing.T, writer Writer) {
    42  	msg := messageGen.NoMux(messageGen.Big)
    43  	msg.MetaData.Store(message.ErrorCode, "200")
    44  	msg.MetaData.Store(message.ErrorMessage, "Hello world!")
    45  	msg.MetaData.Store(message.ErrorMore, "[\"hello world\",123]")
    46  	msg.SetMsgType(message.Return)
    47  	arg := Argument{
    48  		Message: msg,
    49  		Conn:    &NilConn{},
    50  		Encoder: packer.Get("text"),
    51  		Pool: &sync.Pool{
    52  			New: func() interface{} {
    53  				var tmp container.Slice[byte] = make([]byte, mux.MaxBlockSize)
    54  				return &tmp
    55  			},
    56  		},
    57  		OnDebug: nil,
    58  		EHandle: errorhandler.DefaultErrHandler,
    59  	}
    60  	assert.Equal(t, writer.Write(arg, 0), nil)
    61  
    62  	arg.Message.SetMsgType(message.Return)
    63  	assert.Equal(t, writer.Write(arg, 0), nil)
    64  
    65  	// test gzip
    66  	arg.Encoder = packer.Get("gzip")
    67  	assert.Equal(t, writer.Write(arg, 0), nil, "encoder encode failed")
    68  
    69  	arg.Conn = &NilConn{writeFailed: true}
    70  	assert.NotEqual(t, writer.Write(arg, 0), nil, "write return error but Write no return")
    71  }