github.com/cloudwego/kitex@v0.9.0/pkg/remote/trans/gonet/bytebuffer_test.go

/*
 * Copyright 2022 CloudWeGo Authors
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package gonet

import (
	"bufio"
	"bytes"
	"strings"
	"testing"

	"github.com/cloudwego/kitex/internal/test"
	"github.com/cloudwego/kitex/pkg/remote"
)

var (
	msg    = "hello world"
	msgLen = len(msg)
)

// TestBufferReadWrite test bytebuf write and read success.
func TestBufferReadWrite(t *testing.T) {
	var (
		reader = bufio.NewReader(strings.NewReader(strings.Repeat(msg, 5)))
		writer = bufio.NewWriter(bytes.NewBufferString(strings.Repeat(msg, 5)))
		ioRW   = bufio.NewReadWriter(reader, writer)
		bufRW  = NewBufferReadWriter(ioRW)
	)

	testRead(t, bufRW)
	testWrite(t, bufRW)
}

// TestBufferWrite test write success.
func TestBufferWrite(t *testing.T) {
	wi := bytes.NewBufferString(strings.Repeat(msg, 5))
	nbp := &bufferReadWriter{}

	buf := NewBufferWriter(wi)
	testWrite(t, buf)
	testReadFailed(t, buf)

	err := nbp.WriteDirect([]byte(msg), 11)
	// check err not nil
	test.Assert(t, err != nil, err)
}

// TestBufferRead test read success.
func TestBufferRead(t *testing.T) {
	ri := strings.NewReader(strings.Repeat(msg, 5))

	buf := NewBufferReader(ri)
	testWriteFailed(t, buf)
	testRead(t, buf)
}

func testRead(t *testing.T, buf remote.ByteBuffer) {
	var (
		p   []byte
		s   string
		err error
	)

	p, err = buf.Peek(msgLen)
	if err != nil {
		t.Logf("Peek failed, err=%s", err.Error())
		t.FailNow()
	}
	test.Assert(t, string(p) == msg)

	err = buf.Skip(1 + msgLen)
	if err != nil {
		t.Logf("Skip failed, err=%s", err.Error())
		t.FailNow()
	}

	p, err = buf.Next(msgLen - 1)
	if err != nil {
		t.Logf("Next failed, err=%s", err.Error())
		t.FailNow()
	}
	test.Assert(t, string(p) == msg[1:])

	if n := buf.ReadableLen(); n != 3*msgLen {
		t.Logf("ReadableLen expect %d, but got %d", 3*msgLen, n)
		t.FailNow()
	}
	if n := buf.ReadLen(); n != msgLen-1 {
		t.Logf("ReadLen expect %d, but got %d", msgLen-1, n)
		t.FailNow()
	}

	s, err = buf.ReadString(msgLen)
	if err != nil {
		t.Logf("ReadString failed, err=%s", err.Error())
		t.FailNow()
	}
	test.Assert(t, s == msg)

	p, err = buf.ReadBinary(msgLen)
	if err != nil {
		t.Logf("ReadBinary failed, err=%s", err.Error())
		t.FailNow()
	}
	test.Assert(t, string(p) == msg)
}

func testReadFailed(t *testing.T, buf remote.ByteBuffer) {
	p := make([]byte, len(msg))
	var n int

	_, err := buf.Peek(len(msg))
	test.Assert(t, err != nil)

	err = buf.Skip(1)
	test.Assert(t, err != nil)

	_, err = buf.Next(msgLen - 1)
	test.Assert(t, err != nil)

	n = buf.ReadableLen()
	test.Assert(t, n == -1)

	n = buf.ReadLen()
	test.Assert(t, n == 0)

	_, err = buf.ReadString(len(msg))
	test.Assert(t, err != nil)

	_, err = buf.ReadBinary(len(msg))
	test.Assert(t, err != nil)

	n, err = buf.Read(p)
	test.Assert(t, err != nil)
	test.Assert(t, n == -1, n)

	_, err = buf.Read(nil)
	test.Assert(t, err != nil)
}

func testWrite(t *testing.T, buf remote.ByteBuffer) {
	p, err := buf.Malloc(msgLen)
	if err != nil {
		t.Logf("Malloc failed, err=%s", err.Error())
		t.FailNow()
	}
	test.Assert(t, len(p) == msgLen)
	copy(p, msg)

	l := buf.MallocLen()
	test.Assert(t, l == msgLen)

	l, err = buf.WriteString(msg)
	if err != nil {
		t.Logf("WriteString failed, err=%s", err.Error())
		t.FailNow()
	}
	test.Assert(t, l == msgLen)

	l, err = buf.WriteBinary([]byte(msg))
	if err != nil {
		t.Logf("WriteBinary failed, err=%s", err.Error())
		t.FailNow()
	}
	test.Assert(t, l == msgLen)

	err = buf.Flush()
	if err != nil {
		t.Logf("Flush failed, err=%s", err.Error())
		t.FailNow()
	}
}

func testWriteFailed(t *testing.T, buf remote.ByteBuffer) {
	_, err := buf.Malloc(len(msg))
	test.Assert(t, err != nil)

	l := buf.MallocLen()
	test.Assert(t, l == -1)

	_, err = buf.WriteString(msg)
	test.Assert(t, err != nil)

	_, err = buf.WriteBinary([]byte(msg))
	test.Assert(t, err != nil)

	err = buf.Flush()
	test.Assert(t, err != nil)

	n, err := buf.Write([]byte(msg))
	test.Assert(t, err != nil)
	test.Assert(t, n == -1, n)

	_, err = buf.Write(nil)
	test.Assert(t, err != nil)
}