github.com/hasnat/dolt/go@v0.0.0-20210628190320-9eb5d843fbb7/libraries/utils/test/reader.go (about)

     1  // Copyright 2019 Dolthub, Inc.
     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 test
    16  
    17  import (
    18  	"io"
    19  	"math"
    20  
    21  	"github.com/google/uuid"
    22  )
    23  
    24  // TestError is an error that stores a uniquely generated id
    25  type TestError struct {
    26  	ErrId uuid.UUID
    27  }
    28  
    29  // NewTestError creates a TestError with a newly created id
    30  func NewTestError() *TestError {
    31  	id, _ := uuid.NewRandom()
    32  	return &TestError{id}
    33  }
    34  
    35  // Error is the only method defined in the error interface
    36  func (te *TestError) Error() string {
    37  	return "Error! Error! Error!"
    38  }
    39  
    40  // TestReader is an io.Reader that will error after some number of bytes have been read
    41  type TestReader struct {
    42  	data       []byte
    43  	errorAfter int
    44  	totalRead  int
    45  }
    46  
    47  // NewTestReader takes a size, and the number of bytes that should be read before the
    48  // reader errors.  The size is used to generate a buffer of random data that will be
    49  // read by the caller.  After the caller has read the specified number of bytes the
    50  // call to read will fail.
    51  func NewTestReader(dataSize, errorAfter int) *TestReader {
    52  	data := make([]byte, dataSize)
    53  	for i := 0; i < dataSize; i++ {
    54  		data[i] = byte(i % 256)
    55  	}
    56  
    57  	return NewTestReaderWithData(data, errorAfter)
    58  }
    59  
    60  // NewTestReaderWithData creates a TestReader with user supplied data and the number
    61  // of bytes that can be read before errors start occurring on Read calls.
    62  func NewTestReaderWithData(data []byte, errorAfter int) *TestReader {
    63  	return &TestReader{data, errorAfter, 0}
    64  }
    65  
    66  // Read reads data from the internal buffer
    67  func (tr *TestReader) Read(p []byte) (int, error) {
    68  	dataRemaining := len(tr.data) - tr.totalRead
    69  
    70  	if dataRemaining == 0 {
    71  		return 0, io.EOF
    72  	}
    73  
    74  	remainingBeforeErr := math.MaxInt32
    75  	if tr.errorAfter != -1 {
    76  		remainingBeforeErr = tr.errorAfter - tr.totalRead
    77  	}
    78  
    79  	toRead := dataRemaining
    80  	if dataRemaining > remainingBeforeErr {
    81  		toRead = remainingBeforeErr
    82  	}
    83  
    84  	if toRead > len(p) {
    85  		toRead = len(p)
    86  	}
    87  
    88  	var copied int
    89  	if toRead > 0 {
    90  		copied = copy(p[:toRead], tr.data[tr.totalRead:tr.totalRead+toRead])
    91  		tr.totalRead += copied
    92  	}
    93  
    94  	var err error
    95  	if tr.totalRead == tr.errorAfter && copied != len(p) {
    96  		err = NewTestError()
    97  	}
    98  
    99  	return copied, err
   100  }