github.com/cloudwego/hertz@v0.9.3/pkg/network/writer_test.go (about)

     1  /*
     2   * Copyright 2022 CloudWeGo Authors
     3   *
     4   * Licensed under the Apache License, Version 2.0 (the "License");
     5   * you may not use this file except in compliance with the License.
     6   * You may obtain a copy of the License at
     7   *
     8   *     http://www.apache.org/licenses/LICENSE-2.0
     9   *
    10   * Unless required by applicable law or agreed to in writing, software
    11   * distributed under the License is distributed on an "AS IS" BASIS,
    12   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13   * See the License for the specific language governing permissions and
    14   * limitations under the License.
    15   */
    16  
    17  package network
    18  
    19  import (
    20  	"testing"
    21  
    22  	"github.com/cloudwego/hertz/pkg/common/test/assert"
    23  )
    24  
    25  const (
    26  	size1K = 1024
    27  )
    28  
    29  func TestConvertNetworkWriter(t *testing.T) {
    30  	iw := &mockIOWriter{}
    31  	w := NewWriter(iw)
    32  	nw, _ := w.(*networkWriter)
    33  
    34  	// Test malloc
    35  	buf, _ := w.Malloc(size1K)
    36  	assert.DeepEqual(t, len(buf), size1K)
    37  	assert.DeepEqual(t, len(nw.caches), 1)
    38  	assert.DeepEqual(t, len(nw.caches[0].data), size1K)
    39  	assert.DeepEqual(t, cap(nw.caches[0].data), size1K)
    40  	err := w.Flush()
    41  	assert.Nil(t, err)
    42  	assert.DeepEqual(t, size1K, iw.WriteNum)
    43  	assert.DeepEqual(t, len(nw.caches), 0)
    44  	assert.DeepEqual(t, cap(nw.caches), 1)
    45  
    46  	// Test malloc left size
    47  	buf, _ = w.Malloc(size1K + 1)
    48  	assert.DeepEqual(t, len(buf), size1K+1)
    49  	assert.DeepEqual(t, len(nw.caches), 1)
    50  	assert.DeepEqual(t, len(nw.caches[0].data), size1K+1)
    51  	assert.DeepEqual(t, cap(nw.caches[0].data), size1K*2)
    52  	buf, _ = w.Malloc(size1K / 2)
    53  	assert.DeepEqual(t, len(buf), size1K/2)
    54  	assert.DeepEqual(t, len(nw.caches), 1)
    55  	assert.DeepEqual(t, len(nw.caches[0].data), size1K+1+size1K/2)
    56  	assert.DeepEqual(t, cap(nw.caches[0].data), size1K*2)
    57  	buf, _ = w.Malloc(size1K / 2)
    58  	assert.DeepEqual(t, len(buf), size1K/2)
    59  	assert.DeepEqual(t, len(nw.caches), 2)
    60  	assert.DeepEqual(t, len(nw.caches[0].data), size1K+1+size1K/2)
    61  	assert.DeepEqual(t, cap(nw.caches[0].data), size1K*2)
    62  	assert.DeepEqual(t, len(nw.caches[1].data), size1K/2)
    63  	assert.DeepEqual(t, cap(nw.caches[1].data), size1K/2)
    64  	err = w.Flush()
    65  	assert.Nil(t, err)
    66  	assert.DeepEqual(t, size1K*3+1, iw.WriteNum)
    67  	assert.DeepEqual(t, len(nw.caches), 0)
    68  	assert.DeepEqual(t, cap(nw.caches), 2)
    69  
    70  	// Test WriteBinary after Malloc
    71  	buf, _ = w.Malloc(size1K * 6)
    72  	assert.DeepEqual(t, len(buf), size1K*6)
    73  	assert.DeepEqual(t, len(nw.caches[0].data), size1K*6)
    74  	b := make([]byte, size1K)
    75  	w.WriteBinary(b)
    76  	assert.DeepEqual(t, size1K*3+1, iw.WriteNum)
    77  	assert.DeepEqual(t, len(nw.caches[0].data), size1K*7)
    78  	assert.DeepEqual(t, cap(nw.caches[0].data), size1K*8)
    79  
    80  	b = make([]byte, size1K*4)
    81  	w.WriteBinary(b)
    82  	assert.DeepEqual(t, len(nw.caches[1].data), size1K*4)
    83  	assert.DeepEqual(t, cap(nw.caches[1].data), size1K*4)
    84  	assert.DeepEqual(t, nw.caches[1].readOnly, true)
    85  	w.Flush()
    86  	assert.DeepEqual(t, size1K*14+1, iw.WriteNum)
    87  }
    88  
    89  type mockIOWriter struct {
    90  	WriteNum int
    91  }
    92  
    93  func (m *mockIOWriter) Write(p []byte) (n int, err error) {
    94  	m.WriteNum += len(p)
    95  	return len(p), nil
    96  }