github.com/grafana/synthetic-monitoring-agent@v0.64.2/internal/pusher/v2/snappy_concat_test.go (about)

     1  package v2
     2  
     3  import (
     4  	"crypto/rand"
     5  	"io"
     6  	"testing"
     7  
     8  	"github.com/golang/snappy"
     9  	"github.com/stretchr/testify/require"
    10  )
    11  
    12  func TestSnappyConcatReader(t *testing.T) {
    13  	const bigRandSize = 4096 * 4
    14  
    15  	var bigRand [bigRandSize]byte
    16  
    17  	_, err := rand.Read(bigRand[:])
    18  	require.NoError(t, err)
    19  
    20  	for title, tc := range map[string]struct {
    21  		input    [][]byte
    22  		expected []byte
    23  		readFn   func(io.Reader) ([]byte, error)
    24  	}{
    25  		"sample": {
    26  			input: [][]byte{
    27  				snap("HELLO"),
    28  				snap(" WORL"),
    29  				snap("D!"),
    30  			},
    31  			expected: []byte("HELLO WORLD!"),
    32  		},
    33  		"single stream": {
    34  			input: [][]byte{
    35  				snap("Hello World!"),
    36  			},
    37  			expected: []byte("Hello World!"),
    38  		},
    39  		"sample small buf": {
    40  			input: [][]byte{
    41  				snap("HELLO WORLD 1\n"),
    42  				snap("HELLO WORLD 2!\n"),
    43  				snap("HELLO WORLD 3 AND LAST.\n"),
    44  			},
    45  			readFn:   readerWithBufSize(4),
    46  			expected: []byte("HELLO WORLD 1\nHELLO WORLD 2!\nHELLO WORLD 3 AND LAST.\n"),
    47  		},
    48  		"sample tiny buf": {
    49  			input: [][]byte{
    50  				snap("HELLO WORLD 1\n"),
    51  				snap("HELLO WORLD 2!\n"),
    52  				snap("HELLO WORLD 3 AND LAST.\n"),
    53  			},
    54  			readFn:   readerWithBufSize(1),
    55  			expected: []byte("HELLO WORLD 1\nHELLO WORLD 2!\nHELLO WORLD 3 AND LAST.\n"),
    56  		},
    57  		"sample mid buf": {
    58  			input: [][]byte{
    59  				snap("HELLO WORLD 1\n"),
    60  				snap("HELLO WORLD 2!\n"),
    61  				snap("HELLO WORLD 3 AND LAST.\n"),
    62  			},
    63  			readFn:   readerWithBufSize(16),
    64  			expected: []byte("HELLO WORLD 1\nHELLO WORLD 2!\nHELLO WORLD 3 AND LAST.\n"),
    65  		},
    66  		"very large": {
    67  			input: [][]byte{
    68  				snap(string(bigRand[:bigRandSize/4])),
    69  				snap(string(bigRand[bigRandSize/4 : bigRandSize/2])),
    70  				snap(string(bigRand[bigRandSize/2 : 3*bigRandSize/4])),
    71  				snap(string(bigRand[3*bigRandSize/4:])),
    72  			},
    73  			expected: bigRand[:],
    74  		},
    75  		"empty": {},
    76  	} {
    77  		t.Run(title, func(t *testing.T) {
    78  			r := &SnappyConcatReader{
    79  				Streams: tc.input,
    80  			}
    81  			if tc.readFn == nil {
    82  				tc.readFn = io.ReadAll
    83  			}
    84  
    85  			all, err := tc.readFn(r)
    86  			require.NoError(t, err)
    87  
    88  			result, err := snappy.Decode(nil, all)
    89  			require.NoError(t, err)
    90  
    91  			require.Equal(t, tc.expected, result)
    92  		})
    93  	}
    94  }
    95  
    96  func snap(data string) []byte {
    97  	return snappy.Encode(nil, []byte(data))
    98  }
    99  
   100  func readerWithBufSize(n int) func(io.Reader) ([]byte, error) {
   101  	buf := make([]byte, n)
   102  
   103  	return func(r io.Reader) ([]byte, error) {
   104  		var result []byte
   105  
   106  		for {
   107  			n, err := r.Read(buf)
   108  			if err != nil {
   109  				if err == io.EOF {
   110  					return result, nil
   111  				}
   112  
   113  				return nil, err
   114  			}
   115  
   116  			result = append(result, buf[:n]...)
   117  			if n != len(buf) {
   118  				return result, nil
   119  			}
   120  		}
   121  	}
   122  }