go.chromium.org/luci@v0.0.0-20240309015107-7cdc2e660f33/common/data/recordio/writer_test.go (about)

     1  // Copyright 2015 The LUCI Authors.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //      http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package recordio
    16  
    17  import (
    18  	"bytes"
    19  	"encoding/binary"
    20  	"errors"
    21  	"testing"
    22  
    23  	. "github.com/smartystreets/goconvey/convey"
    24  )
    25  
    26  type testWriter struct {
    27  	buf  bytes.Buffer
    28  	errQ []error
    29  }
    30  
    31  func (w *testWriter) Write(data []byte) (int, error) {
    32  	if len(w.errQ) > 0 {
    33  		err := w.errQ[0]
    34  		w.errQ = w.errQ[1:]
    35  
    36  		if err != nil {
    37  			return 0, err
    38  		}
    39  	}
    40  	return w.buf.Write(data)
    41  }
    42  
    43  func TestWriteFrame(t *testing.T) {
    44  	t.Parallel()
    45  
    46  	Convey(`Using a buffered test writer`, t, func() {
    47  		tw := &testWriter{}
    48  
    49  		Convey(`WriteFrame will successfully encode a zero-byte frame.`, func() {
    50  			count, err := WriteFrame(tw, []byte{})
    51  			So(count, ShouldEqual, 1)
    52  			So(err, ShouldEqual, nil)
    53  			So(tw.buf.Bytes(), ShouldResemble, []byte{0x00})
    54  		})
    55  
    56  		Convey(`WriteFrame will successfully encode a 129-byte frame.`, func() {
    57  			data := bytes.Repeat([]byte{0x5A}, 129)
    58  
    59  			count, err := WriteFrame(tw, data)
    60  			So(count, ShouldEqual, 131)
    61  			So(err, ShouldEqual, nil)
    62  			So(tw.buf.Bytes(), ShouldResemble, append([]byte{0x81, 0x01}, data...))
    63  		})
    64  
    65  		Convey(`Will return an error when failing to write the size header.`, func() {
    66  			failErr := errors.New("test: test-induced size error")
    67  			tw.errQ = []error{
    68  				failErr,
    69  			}
    70  
    71  			count, err := WriteFrame(tw, []byte{0xd0, 0x65})
    72  			So(count, ShouldEqual, 0)
    73  			So(err, ShouldEqual, failErr)
    74  		})
    75  
    76  		Convey(`Will return an error when failing to write the frame data.`, func() {
    77  			failErr := errors.New("test: test-induced size error")
    78  			tw.errQ = []error{
    79  				nil,
    80  				failErr,
    81  			}
    82  
    83  			count, err := WriteFrame(tw, []byte{0xd0, 0x65})
    84  			So(count, ShouldEqual, 1)
    85  			So(err, ShouldEqual, failErr)
    86  		})
    87  	})
    88  }
    89  
    90  func TestWriter(t *testing.T) {
    91  	t.Parallel()
    92  
    93  	Convey(`A Writer, configured to write to a buffer`, t, func() {
    94  		tw := &testWriter{}
    95  		w := NewWriter(tw)
    96  
    97  		Convey(`Can write consecutive frames in 3-byte chunks.`, func() {
    98  			expected := []byte{}
    99  			var sizeBuf [binary.MaxVarintLen64]byte
   100  			for _, size := range []int{
   101  				1,
   102  				0,
   103  				1025,
   104  				129,
   105  				11,
   106  			} {
   107  				b := bytes.Repeat([]byte{0x55}, size)
   108  				expected = append(expected, sizeBuf[:binary.PutUvarint(sizeBuf[:], uint64(size))]...)
   109  				expected = append(expected, b...)
   110  
   111  				for len(b) > 0 {
   112  					count := 3
   113  					if count > len(b) {
   114  						count = len(b)
   115  					}
   116  
   117  					c, err := w.Write(b[:count])
   118  					So(err, ShouldBeNil)
   119  					So(c, ShouldEqual, count)
   120  
   121  					b = b[count:]
   122  				}
   123  				So(w.Flush(), ShouldBeNil)
   124  			}
   125  
   126  			So(tw.buf.Bytes(), ShouldResemble, expected)
   127  		})
   128  
   129  		Convey(`Will write empty frames if Flush()ed.`, func() {
   130  			So(w.Flush(), ShouldBeNil)
   131  			So(w.Flush(), ShouldBeNil)
   132  			So(w.Flush(), ShouldBeNil)
   133  			So(tw.buf.Bytes(), ShouldResemble, []byte{0x00, 0x00, 0x00})
   134  		})
   135  
   136  		Convey(`Will fail to Flush() if the Write fails.`, func() {
   137  			failErr := errors.New("test: test-induced size error")
   138  			tw.errQ = []error{
   139  				failErr,
   140  			}
   141  			So(w.Flush(), ShouldEqual, failErr)
   142  		})
   143  	})
   144  }