github.com/Finschia/finschia-sdk@v0.48.1/snapshots/chunk_test.go (about)

     1  package snapshots_test
     2  
     3  import (
     4  	"bytes"
     5  	"errors"
     6  	"io"
     7  	"testing"
     8  
     9  	"github.com/stretchr/testify/assert"
    10  	"github.com/stretchr/testify/require"
    11  
    12  	"github.com/Finschia/finschia-sdk/snapshots"
    13  )
    14  
    15  func TestChunkWriter(t *testing.T) {
    16  	ch := make(chan io.ReadCloser, 100)
    17  	go func() {
    18  		chunkWriter := snapshots.NewChunkWriter(ch, 2)
    19  
    20  		n, err := chunkWriter.Write([]byte{1, 2, 3})
    21  		require.NoError(t, err)
    22  		assert.Equal(t, 3, n)
    23  
    24  		n, err = chunkWriter.Write([]byte{4, 5, 6})
    25  		require.NoError(t, err)
    26  		assert.Equal(t, 3, n)
    27  
    28  		n, err = chunkWriter.Write([]byte{7, 8, 9})
    29  		require.NoError(t, err)
    30  		assert.Equal(t, 3, n)
    31  
    32  		err = chunkWriter.Close()
    33  		require.NoError(t, err)
    34  
    35  		// closed writer should error
    36  		_, err = chunkWriter.Write([]byte{10})
    37  		require.Error(t, err)
    38  
    39  		// closing again should be fine
    40  		err = chunkWriter.Close()
    41  		require.NoError(t, err)
    42  	}()
    43  
    44  	assert.Equal(t, [][]byte{{1, 2}, {3, 4}, {5, 6}, {7, 8}, {9}}, readChunks(ch))
    45  
    46  	// 0-sized chunks should return the whole body as one chunk
    47  	ch = make(chan io.ReadCloser, 100)
    48  	go func() {
    49  		chunkWriter := snapshots.NewChunkWriter(ch, 0)
    50  		_, err := chunkWriter.Write([]byte{1, 2, 3})
    51  		require.NoError(t, err)
    52  		_, err = chunkWriter.Write([]byte{4, 5, 6})
    53  		require.NoError(t, err)
    54  		err = chunkWriter.Close()
    55  		require.NoError(t, err)
    56  	}()
    57  	assert.Equal(t, [][]byte{{1, 2, 3, 4, 5, 6}}, readChunks(ch))
    58  
    59  	// closing with error should return the error
    60  	theErr := errors.New("boom")
    61  	ch = make(chan io.ReadCloser, 100)
    62  	go func() {
    63  		chunkWriter := snapshots.NewChunkWriter(ch, 2)
    64  		_, err := chunkWriter.Write([]byte{1, 2, 3})
    65  		require.NoError(t, err)
    66  		chunkWriter.CloseWithError(theErr)
    67  	}()
    68  	chunk, err := io.ReadAll(<-ch)
    69  	require.NoError(t, err)
    70  	assert.Equal(t, []byte{1, 2}, chunk)
    71  	_, err = io.ReadAll(<-ch)
    72  	require.Error(t, err)
    73  	assert.Equal(t, theErr, err)
    74  	assert.Empty(t, ch)
    75  
    76  	// closing immediately should return no chunks
    77  	ch = make(chan io.ReadCloser, 100)
    78  	chunkWriter := snapshots.NewChunkWriter(ch, 2)
    79  	err = chunkWriter.Close()
    80  	require.NoError(t, err)
    81  	assert.Empty(t, ch)
    82  }
    83  
    84  func TestChunkReader(t *testing.T) {
    85  	ch := makeChunks([][]byte{
    86  		{1, 2, 3},
    87  		{4},
    88  		{},
    89  		{5, 6},
    90  	})
    91  	chunkReader := snapshots.NewChunkReader(ch)
    92  
    93  	buf := []byte{0, 0, 0, 0}
    94  	n, err := chunkReader.Read(buf)
    95  	require.NoError(t, err)
    96  	assert.Equal(t, 3, n)
    97  	assert.Equal(t, []byte{1, 2, 3, 0}, buf)
    98  
    99  	buf = []byte{0, 0, 0, 0}
   100  	n, err = chunkReader.Read(buf)
   101  	require.NoError(t, err)
   102  	assert.Equal(t, 1, n)
   103  	assert.Equal(t, []byte{4, 0, 0, 0}, buf)
   104  
   105  	buf = []byte{0, 0, 0, 0}
   106  	n, err = chunkReader.Read(buf)
   107  	require.NoError(t, err)
   108  	assert.Equal(t, 2, n)
   109  	assert.Equal(t, []byte{5, 6, 0, 0}, buf)
   110  
   111  	buf = []byte{0, 0, 0, 0}
   112  	_, err = chunkReader.Read(buf)
   113  	require.Error(t, err)
   114  	assert.Equal(t, io.EOF, err)
   115  
   116  	err = chunkReader.Close()
   117  	require.NoError(t, err)
   118  
   119  	err = chunkReader.Close() // closing twice should be fine
   120  	require.NoError(t, err)
   121  
   122  	// Empty channel should be fine
   123  	ch = makeChunks(nil)
   124  	chunkReader = snapshots.NewChunkReader(ch)
   125  	buf = make([]byte, 4)
   126  	_, err = chunkReader.Read(buf)
   127  	require.Error(t, err)
   128  	assert.Equal(t, io.EOF, err)
   129  
   130  	// Using a pipe that closes with an error should return the error
   131  	theErr := errors.New("boom")
   132  	pr, pw := io.Pipe()
   133  	pch := make(chan io.ReadCloser, 1)
   134  	pch <- pr
   135  	pw.CloseWithError(theErr)
   136  
   137  	chunkReader = snapshots.NewChunkReader(pch)
   138  	buf = make([]byte, 4)
   139  	_, err = chunkReader.Read(buf)
   140  	require.Error(t, err)
   141  	assert.Equal(t, theErr, err)
   142  
   143  	// Closing the reader should close the writer
   144  	pr, pw = io.Pipe()
   145  	pch = make(chan io.ReadCloser, 2)
   146  	pch <- io.NopCloser(bytes.NewBuffer([]byte{1, 2, 3}))
   147  	pch <- pr
   148  	close(pch)
   149  
   150  	go func() {
   151  		chunkReader := snapshots.NewChunkReader(pch)
   152  		buf := []byte{0, 0, 0, 0}
   153  		_, err := chunkReader.Read(buf)
   154  		require.NoError(t, err)
   155  		assert.Equal(t, []byte{1, 2, 3, 0}, buf)
   156  
   157  		err = chunkReader.Close()
   158  		require.NoError(t, err)
   159  	}()
   160  
   161  	_, err = pw.Write([]byte{9, 9, 9})
   162  	require.Error(t, err)
   163  	assert.Equal(t, err, io.ErrClosedPipe)
   164  }