github.com/Schaudge/grailbase@v0.0.0-20240223061707-44c758a471c0/recordio/recordioutil/v2_test.go (about)

     1  // Copyright 2018 GRAIL, Inc. All rights reserved.
     2  // Use of this source code is governed by the Apache-2.0
     3  // license that can be found in the LICENSE file.
     4  
     5  package recordioutil_test
     6  
     7  import (
     8  	"bytes"
     9  	"testing"
    10  
    11  	"github.com/Schaudge/grailbase/fileio"
    12  	"github.com/Schaudge/grailbase/recordio"
    13  	"github.com/Schaudge/grailbase/recordio/deprecated"
    14  	"github.com/Schaudge/grailbase/recordio/recordioutil"
    15  	"github.com/stretchr/testify/require"
    16  )
    17  
    18  func readV1(t *testing.T, format fileio.FileType, buf *bytes.Buffer) (s []string) {
    19  	opts, err := recordioutil.ScannerOptsFromName("foo." + fileio.FileSuffix(format))
    20  	if err != nil {
    21  		s = append(s, err.Error())
    22  	}
    23  	sc := recordio.NewScanner(bytes.NewReader(buf.Bytes()), opts)
    24  	for sc.Scan() {
    25  		s = append(s, string(sc.Get().([]byte)))
    26  	}
    27  	if err := sc.Err(); err != nil {
    28  		s = append(s, err.Error())
    29  	}
    30  	return
    31  }
    32  
    33  func TestPacked(t *testing.T) {
    34  	buf := &bytes.Buffer{}
    35  	w := deprecated.NewLegacyPackedWriter(buf, deprecated.LegacyPackedWriterOpts{})
    36  	_, err := w.Write([]byte("Foo"))
    37  	require.NoError(t, err)
    38  	_, err = w.Write([]byte("Baz"))
    39  	require.NoError(t, err)
    40  	w.Flush()
    41  	require.Equal(t, []string{"Foo", "Baz"}, readV1(t, fileio.GrailRIOPacked, buf))
    42  }
    43  
    44  func TestUnpacked(t *testing.T) {
    45  	buf := &bytes.Buffer{}
    46  	w := deprecated.NewLegacyWriter(buf, deprecated.LegacyWriterOpts{})
    47  	_, err := w.Write([]byte("Foo"))
    48  	require.NoError(t, err)
    49  	_, err = w.Write([]byte("Baz"))
    50  	require.NoError(t, err)
    51  	require.Equal(t, []string{"Foo", "Baz"}, readV1(t, fileio.GrailRIO, buf))
    52  }
    53  
    54  func TestCompressed(t *testing.T) {
    55  	buf := &bytes.Buffer{}
    56  	w := deprecated.NewLegacyPackedWriter(buf, deprecated.LegacyPackedWriterOpts{
    57  		Transform: recordioutil.NewFlateTransform(-1).CompressTransform})
    58  	_, err := w.Write([]byte("Foo"))
    59  	require.NoError(t, err)
    60  	_, err = w.Write([]byte("Baz"))
    61  	require.NoError(t, err)
    62  	w.Flush()
    63  	require.Equal(t, []string{"Foo", "Baz"}, readV1(t, fileio.GrailRIOPackedCompressed, buf))
    64  }