github.com/blend/go-sdk@v1.20220411.3/sh/limit_bytes_test.go (about)

     1  /*
     2  
     3  Copyright (c) 2022 - Present. Blend Labs, Inc. All rights reserved
     4  Use of this source code is governed by a MIT license that can be found in the LICENSE file.
     5  
     6  */
     7  
     8  package sh
     9  
    10  import (
    11  	"bytes"
    12  	"math/rand"
    13  	"testing"
    14  	"time"
    15  
    16  	"github.com/blend/go-sdk/assert"
    17  	"github.com/blend/go-sdk/ex"
    18  )
    19  
    20  func TestMaxBytesWriter(t *testing.T) {
    21  	assert := assert.New(t)
    22  
    23  	buf := bytes.NewBuffer(nil)
    24  	mbw := LimitBytes(64, buf)
    25  
    26  	written, err := mbw.Write(makeChunk(32))
    27  	assert.Nil(err)
    28  	assert.Equal(32, written)
    29  
    30  	written, err = mbw.Write(makeChunk(16))
    31  	assert.Nil(err)
    32  	assert.Equal(16, written)
    33  
    34  	written, err = mbw.Write(makeChunk(32))
    35  	assert.True(ex.Is(err, ErrMaxBytesWriterCapacityLimit))
    36  	assert.Equal(0, written)
    37  }
    38  
    39  func makeChunk(len int) []byte {
    40  	output := make([]byte, len)
    41  	for x := 0; x < len; x++ {
    42  		output[x] = randomLetter()
    43  	}
    44  	return output
    45  }
    46  
    47  var (
    48  	provider = rand.New(rand.NewSource(time.Now().UnixNano()))
    49  
    50  	// LowerLetters is a runset of lowercase letters.
    51  	lowerLetters = []byte("abcdefghijklmnopqrstuvwxyz")
    52  
    53  	// UpperLetters is a runset of uppercase letters.
    54  	upperLetters = []byte("ABCDEFGHIJKLMNOPQRSTUVWXYZ")
    55  
    56  	// Letters is a runset of both lower and uppercase letters.
    57  	letters = append(lowerLetters, upperLetters...)
    58  )
    59  
    60  func randomLetter() byte {
    61  	return letters[provider.Intn(len(letters))]
    62  }